[Unity][SLM] GPU sampling (#16575)

This PR adds GPU sampling support to SLM
apache · Feb 23, 2024 · 84b3f69 · 84b3f69
1 parent faa6628
commit 84b3f69
Show file tree

Hide file tree

Showing 5 changed files with 973 additions and 3 deletions.
diff --git a/python/tvm/relax/frontend/nn/_tensor_op.py b/python/tvm/relax/frontend/nn/_tensor_op.py
@@ -67,6 +67,22 @@ def __truediv__(self, other):
         other = _convert_scalar(other, self)
         return _op().divide(self, other)
 
+    def __lt__(self, other):
+        other = _convert_scalar(other, self)
+        return _op().less(self, other)
+
+    def __le__(self, other):
+        other = _convert_scalar(other, self)
+        return _op().less_equal(self, other)
+
+    def __gt__(self, other):
+        other = _convert_scalar(other, self)
+        return _op().greater(self, other)
+
+    def __ge__(self, other):
+        other = _convert_scalar(other, self)
+        return _op().greater_equal(self, other)
+
     def astype(self, dtype):
         return _op().astype(self, dtype)