modified conditions for parameters

Author: Pratheesh-04-MCW

python/tvm/relax/frontend/torch/base_fx_graph_translator.py

@@ -1148,21 +1148,21 @@ def _gather(self, node: fx.Node) -> relax.Var:
         index = self.env[node.args[2]]
         return self.block_builder.emit(relax.op.gather_elements(x, index, axis=dim))
 
-    def _index_put_(self, node: fx.Node) -> relax.Var:
+    def _index_put(self, node: fx.Node) -> relax.Var:
         args = self.retrieve_args(node)
         tensor = args[0]
-        indices = args[1] if len(args) > 1 else node.kwargs.get("indices", ())
+        indices = args[1] if len(args) > 1 else node.kwargs.get("indices")
         values = args[2] if len(args) > 2 else node.kwargs.get("values")
         accumulate = args[3] if len(args) > 3 else node.kwargs.get("accumulate", False)
 
-        # Ensure accumulate is a boolean
-        if isinstance(accumulate, str):
-            accumulate = accumulate.lower() == "true"
-        elif not isinstance(accumulate, bool):
-            accumulate = bool(accumulate)
+        if indices is None or values is None:
+            raise ValueError("'indices and values' arguments are required for index_put operation")
+        
+        if not isinstance(accumulate, bool):
+            raise TypeError("'accumulate' must be a boolean value, got {}".format(type(accumulate)))
 
         if isinstance(indices, (list, tuple)):
-            indices = relax.Tuple(indices) if indices else relax.Tuple([])
+            indices = relax.Tuple(indices)
         return self.block_builder.emit(relax.op.index_put(tensor, indices, values, accumulate))
 
     def _permute(self, node: fx.Node) -> relax.Var:

python/tvm/relax/frontend/torch/exported_program_translator.py

@@ -420,7 +420,7 @@ def create_convert_map(
             "flatten.using_ints": self._flatten,
             "flip.default": self._flip,
             "gather.default": self._gather,
-            "index_put_.default": self._index_put_,
+            "index_put_.default": self._index_put,
             "narrow.default": self._narrow,
             "permute.default": self._permute,
             "repeat.default": self._repeat,

python/tvm/relax/frontend/torch/fx_translator.py

@@ -746,7 +746,7 @@ def create_convert_map(
             "flatten": self._flatten,
             "flip": self._flip,
             "gather": self._gather,
-            "index_put_": self._index_put_,
+            "index_put_": self._index_put,
             "narrow": self._narrow,
             "numel": self._numel,
             "permute": self._permute,

python/tvm/relax/op/manipulate.py

@@ -588,7 +588,7 @@ def index_put(
         ]
     """
     if not isinstance(indices, (list, tuple)):
-        indices = RxTuple(indices) if indices else RxTuple([])
+        indices = RxTuple(indices)
     return _ffi_api.index_put(data, indices, values, accumulate)  # type: ignore
 
 

python/tvm/topi/index_put.py

@@ -18,6 +18,7 @@
 from tvm import te
 from tvm import tir
 from . import utils
+import math
 
 
 def index_put(data, indices, values, accumulate=False):
@@ -56,14 +57,12 @@ def index_put(data, indices, values, accumulate=False):
 
     # Prepare ranges and strides
     shape = data.shape
-    full_range = 1
-    for dim in shape:
-        full_range *= dim
+    full_range = math.prod(data.shape)
 
     # Check all indices have same length
-    index_len = indices[0].shape[0]
+    index_len = len(indices[0])
     for idx in indices[1:]:
-        if not utils.equal_const_int(idx.shape[0], index_len):
+        if not utils.equal_const_int(len(idx), index_len):
             raise ValueError("All index tensors must have same length")
 
     def gen_ir(data_ptr, index_ptrs, values_ptr, out_ptr, reduce_func):