Setitem with tensor values. And Boolean type promotion (#290)

zhumakhan · Zhumakhan · vadiklyutiy · commit 60e75ca41ffe · 2024-07-23T21:48:33.000+04:00
- setitem with tensor values
- boolean type promotion

---------

Co-authored-by: Zhumakhan &lt;nazirzhumakhan@gmail,.com&gt;
diff --git a/python/hidet/graph/frontend/torch/register_functions.py b/python/hidet/graph/frontend/torch/register_functions.py
@@ -269,9 +269,6 @@ def setitem(x: Tensor, item, setvalue):
     if not isinstance(item, tuple):
         item = tuple([item])
 
-    if not isinstance(setvalue, (int, float)):
-        raise NotImplementedError('Currently Tensor __setitem__ only supports int or float values')
-
     # now, the item could have
     # 1. integer index
     # 2. slice
@@ -321,7 +318,11 @@ def setitem(x: Tensor, item, setvalue):
             ends.append(v.stop)
             steps.append(v.step)
 
-    out = ops.set_strided_slice(x, starts, ends, steps, setvalue)
+    if isinstance(setvalue, Tensor):
+        squeeze_dims = [i for i, dimlen in enumerate(setvalue.shape) if dimlen == 1]
+        setvalue = ops.squeeze(setvalue, squeeze_dims)
+
+    out = ops.set_strided_slice(x, setvalue, starts, ends, steps)
     return out
 
 
diff --git a/python/hidet/graph/ops/arithmetic.py b/python/hidet/graph/ops/arithmetic.py
@@ -192,11 +192,11 @@ class SetStridedSliceTask(Task):
     def __init__(
         self,
         data: TensorNode,
+        setvalue: Union[int, float, TensorNode],
         starts: List[Optional[int]],
         ends: List[Optional[int]],
         axes: List[int],
         strides: List[int],
-        setvalue: [Union[int, float]],
     ):
         assert len(starts) == len(ends) == len(axes) == len(strides)
         if len(axes) != len(set(axes)):
@@ -216,17 +216,37 @@ def __init__(
                 )
             axis2info[axis] = (start, end, stride)
 
+        if isinstance(setvalue, TensorNode):
+            inputs = [data, setvalue]
+        else:
+            inputs = [data]
+
         def fmap(indices):
-            ret = data.type.dtype(setvalue)
+            if isinstance(setvalue, TensorNode):
+                ret = setvalue
+            else:
+                ret = data.type.dtype(setvalue)
+
+            new_val = True
+            new_indices = []
             for axis, index in enumerate(indices):
                 start, end, stride = axis2info[axis]
-                ret = if_then_else(
-                    logical_or(index < start, index >= end, (index - start) % stride != 0), data[indices], ret
+                new_val = if_then_else(
+                    logical_or(index < start, index >= end, (index - start) % stride != 0), False, new_val
                 )
-            return ret
+                if start + 1 < end:
+                    new_indices.append((index - start) // stride)
+
+            if isinstance(setvalue, TensorNode):
+                if len(new_indices) != 0:
+                    ret = ret[new_indices]
+                else:
+                    ret = ret[(0,) * ret.ndim]
+
+            return if_then_else(new_val, ret, data[indices])
 
         out = compute('out', shape=output_shape, fcompute=lambda *indices: fmap(indices))
-        super().__init__(name='set_slice', inputs=[data], outputs=[out])
+        super().__init__(name='set_slice', inputs=inputs, outputs=[out])
 
 
 class RollTask(Task):
@@ -702,18 +722,29 @@ class SetStridedSliceOp(Operator):
     def __init__(
         self,
         data: Tensor,
+        setvalue: Optional[Union[int, float, Tensor]],
         starts: Sequence[Optional[int]],
         ends: Sequence[Optional[int]],
         strides: Optional[Sequence[Optional[int]]] = None,
-        setvalue: Optional[Union[int, float]] = 0.0,
     ):
+        if setvalue is None:
+            setvalue = 0.0
+
         starts, ends, axes, strides = normalize_slice(data.shape, starts, ends, axes=None, strides=strides)
-        task = SetStridedSliceTask(input_like(data, 'data'), starts, ends, axes, strides, setvalue)
-        super().__init__(
-            inputs=[data],
-            attributes={'starts': starts, 'ends': ends, 'strides': strides, 'setvalue': setvalue},
-            task=task,
-        )
+
+        attributes = {'starts': starts, 'ends': ends, 'strides': strides}
+
+        if isinstance(setvalue, Tensor):
+            task = SetStridedSliceTask(
+                input_like(data, 'data'), input_like(setvalue, 'setvalue'), starts, ends, axes, strides
+            )
+            inputs = [data, setvalue]
+        else:
+            task = SetStridedSliceTask(input_like(data, 'data'), setvalue, starts, ends, axes, strides)
+            inputs = [data]
+            attributes['setvalue'] = setvalue
+
+        super().__init__(inputs=inputs, attributes=attributes, task=task)
 
 
 class RollOp(Operator):
@@ -1070,9 +1101,9 @@ def composite_elementwise(
 
 def set_strided_slice(
     data: Tensor,
+    setvalue: Optional[Union[int, float]],
     starts: Sequence[Optional[int]],
     ends: Sequence[Optional[int]],
     strides: Optional[Sequence[Optional[int]]] = None,
-    setvalue: Optional[Union[int, float]] = 0.0,
 ) -> Tensor:
-    return SetStridedSliceOp(data, starts, ends, strides, setvalue).outputs[0]
+    return SetStridedSliceOp(data, setvalue, starts, ends, strides).outputs[0]
diff --git a/python/hidet/ir/dtypes/boolean.py b/python/hidet/ir/dtypes/boolean.py
@@ -35,6 +35,9 @@ def is_complex(self) -> bool:
     def is_vector(self) -> bool:
         return False
 
+    def is_boolean(self) -> bool:
+        return True
+
     def constant(self, value: Any):
         from hidet.ir.expr import constant
 
diff --git a/python/hidet/ir/dtypes/complex.py b/python/hidet/ir/dtypes/complex.py
@@ -35,6 +35,9 @@ def is_vector(self) -> bool:
     def is_complex(self) -> bool:
         return True
 
+    def is_boolean(self) -> bool:
+        return False
+
     def constant(self, value: Any):
         from hidet.ir.expr import Constant, constant
 
diff --git a/python/hidet/ir/dtypes/floats.py b/python/hidet/ir/dtypes/floats.py
@@ -51,6 +51,9 @@ def is_complex(self) -> bool:
     def is_vector(self) -> bool:
         return False
 
+    def is_boolean(self) -> bool:
+        return False
+
     def constant(self, value: Any):
         from hidet.ir.expr import Constant, constant
 
diff --git a/python/hidet/ir/dtypes/integer.py b/python/hidet/ir/dtypes/integer.py
@@ -45,6 +45,9 @@ def is_complex(self) -> bool:
     def is_vector(self) -> bool:
         return False
 
+    def is_boolean(self) -> bool:
+        return False
+
     def constant(self, value: Any):
         from hidet.ir.expr import Constant, constant
 
diff --git a/python/hidet/ir/dtypes/promotion.py b/python/hidet/ir/dtypes/promotion.py
@@ -131,6 +131,8 @@ def promote_type(t1: DataType, t2: DataType) -> DataType:
         return t2
     elif t1.is_float() and t2.is_integer():
         return t1
+    elif t1.is_boolean() and t2.is_boolean():
+        return t1
     else:
         pair = (t1, t2)
         if pair not in _promotion_table:
diff --git a/python/hidet/ir/type.py b/python/hidet/ir/type.py
@@ -174,6 +174,9 @@ def is_complex(self) -> bool:
     def is_vector(self) -> bool:
         raise NotImplementedError()
 
+    def is_boolean(self) -> bool:
+        raise NotImplementedError()
+
     def constant(self, value: Any):
         raise NotImplementedError()
 
diff --git a/tests/frontends/torch/test_torch_inplace.py b/tests/frontends/torch/test_torch_inplace.py
@@ -0,0 +1,64 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import pytest
+import torch
+from hidet.testing.torch_utils import check_module, FunctionalModule
+
+
+@pytest.mark.parametrize(
+    "a_shape, b_shape, indices",
+    [
+        [[10, 11, 12, 13], [10, 9, 8], (slice(None), slice(9), slice(8), 0)],
+        [[5, 4, 3, 2], [2, 2], (slice(2, 4), slice(2, 4), 0, 1)],
+        [[1, 1, 1024], [10], (0, 0, slice(1000, 1010))],
+        [[4, 4, 4, 4], [1, 1, 1], (slice(1), 0, slice(1), slice(1))],
+        [[4, 4, 4, 4], [2, 3], (3, slice(2), 2, slice(3))],
+        [[4, 4, 4, 4], [2, 3], (slice(2), 0, slice(3), 0)],
+        [[4, 4, 4, 4], [4, 4], (0, Ellipsis, 0)],
+        [[10, 10], [10, 10], (Ellipsis,)],
+        [[1, 3, 28, 28, 85], [1, 3, 28, 28, 2], (Ellipsis, slice(2))],
+    ],
+)
+def test_setitem_with_tensor(a_shape, b_shape, indices):
+    def check_setitem(x, y, indices):
+        x[indices] = y
+        return x
+
+    check_module(
+        FunctionalModule(op=check_setitem), args=[torch.randn(a_shape), torch.randn(b_shape), indices], atol=0, rtol=0
+    )
+
+
+@pytest.mark.parametrize(
+    "a_shape, setvalue, indices",
+    [
+        [[10, 11, 12, 13], 1.0, (slice(None), slice(9), slice(8), 0)],
+        [[5, 4, 3, 2], 1.0, (slice(2, 4), slice(2, 4), 0, 1)],
+        [[1, 1, 1024], 1.0, (0, 0, slice(1000, 1010))],
+        [[4, 4, 4, 4], 1.0, (slice(1), 0, slice(1), slice(1))],
+        [[4, 4, 4, 4], 1.0, (3, slice(2), 2, slice(3))],
+        [[4, 4, 4, 4], 1.0, (slice(2), 0, slice(3), 0)],
+        [[4, 4, 4, 4], 1.0, (0, Ellipsis, 0)],
+        [[10, 10], 1.0, (Ellipsis,)],
+        [[1, 3, 28, 28, 85], 1.0, (Ellipsis, slice(2))],
+    ],
+)
+def test_setitem_with_scalar(a_shape, setvalue, indices):
+    def check_setitem(x, setvalue, indices):
+        x[indices] = setvalue
+        return x
+
+    check_module(FunctionalModule(op=check_setitem), args=[torch.randn(a_shape), setvalue, indices], atol=0, rtol=0)
+
+
+if __name__ == '__main__':
+    pytest.main([__file__])
