Skip to content

Commit fac6ab2

Browse files
author
Ramana Radhakrishnan
committed
Add support for overloading comparison operations in relay (#2910)
This commit adds support for overloaded comparison operators in the python binding for relay. Add a testcase for this
1 parent 1fdf111 commit fac6ab2

File tree

2 files changed

+68
-0
lines changed

2 files changed

+68
-0
lines changed

python/tvm/relay/expr.py

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,38 @@ def astype(self, dtype):
7070
def __neg__(self):
7171
return _op_make.negative(self)
7272

73+
def __lt__(self, other):
74+
if isinstance(other, Expr):
75+
return _op_make.less(self, other)
76+
elif isinstance(other, _Number):
77+
raise TypeError('convert "%s" with `const` first' % str(other))
78+
else:
79+
raise TypeError("type %s not supported" % str(type(other)))
80+
81+
def __gt__(self, other):
82+
if isinstance(other, Expr):
83+
return _op_make.greater(self, other)
84+
elif isinstance(other, _Number):
85+
raise TypeError('convert "%s" with `const` first' % str(other))
86+
else:
87+
raise TypeError("type %s not supported" % str(type(other)))
88+
89+
def __ge__(self, other):
90+
if isinstance(other, Expr):
91+
return _op_make.greater_equal(self, other)
92+
elif isinstance(other, _Number):
93+
raise TypeError('convert "%s" with `const` first' % str(other))
94+
else:
95+
raise TypeError("type %s not supported" % str(type(other)))
96+
97+
def __le__(self, other):
98+
if isinstance(other, Expr):
99+
return _op_make.less_equal(self, other)
100+
elif isinstance(other, _Number):
101+
raise TypeError('convert "%s" with `const` first' % str(other))
102+
else:
103+
raise TypeError("type %s not supported" % str(type(other)))
104+
73105
def __add__(self, other):
74106
if isinstance(other, Expr):
75107
return _op_make.add(self, other)

tests/python/relay/test_cmp_op.py

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
18+
from tvm import relay
19+
a = relay.Var("a")
20+
b = relay.expr.const (1.0, dtype='float32')
21+
22+
c = a < b
23+
d = relay.less (a, b)
24+
assert (c.astext() == d.astext())
25+
26+
c = a > b
27+
d = relay.greater (a, b)
28+
assert (c.astext() == d.astext())
29+
30+
c = (a >= b)
31+
d = relay.greater_equal(a, b)
32+
assert (c.astext() == d.astext())
33+
34+
c = (a <= b)
35+
d = relay.less_equal(a, b)
36+
assert (c.astext() == d.astext())

0 commit comments

Comments
 (0)