fix(index_put): handle None cases

Author: AyoubMDL

onnxscript/function_libs/torch_lib/ops/core.py

@@ -1647,7 +1647,8 @@ def aten_chunk(self: TTensor, chunks: int, dim: int = 0) -> Sequence[TTensor]:
     dim_size = op.Gather(self_shape, dim, axis=0)
     # Compute size/chunk to get the number of data in one chunk
     num_per_chunk = op.Div(dim_size, chunks)
-    num_per_chunk = op.Cast(op.Mod(dim_size, chunks) > 0, to=INT64.dtype) + num_per_chunk  # type: ignore[operator]
+    num_per_chunk = op.Cast(op.Mod(dim_size, chunks) > 0, to=INT64.dtype) + \
+        num_per_chunk  # type: ignore[operator]
 
     # Compute real chunk number
     num_chunk = op.Div(dim_size, num_per_chunk)
@@ -4259,15 +4260,51 @@ def aten_index_put(
     See implementation of `torch.onnx.symbolic_opset11.index_put
     <https://github.com/pytorch/pytorch/blob/main/torch/onnx/symbolic_opset11.py#L212>`_.
     """
+    # Pad indices with None so It has the same rank as self
+    self_rank = len(self.shape)
+    if len(indices) < self_rank:
+        indices = list(indices) + [None] * (self_rank - len(indices))
+
+    values_shape = values.shape.numpy()
+    # Pad values_shape with 1 so It has the same rank as self
+    if len(values_shape) < self_rank:
+        values_shape = (1,) * (self_rank - len(values_shape)) + values_shape
+
+    index_vectors = []
+    for i, index in enumerate(indices):
+        if index is None:
+            # For a full slice, create a range.
+            index_vector = op.Range(start=0, limit=values_shape[i], delta=1)
+        else:
+            index_vector = index
 
-    # TODO(justinchuby): Handle when indicies has more than one element
-    index = indices[0]
-    new_index = op.Unsqueeze(index, [-1])
+        # Shape vector with 1s, except at axis i.
+        shape_vector = [1] * self_rank
+        shape_vector[i] = values_shape[i]
+
+        # Reshape index_vector so that only the i-th dimension matches values_shape[i]
+        reshaped_index_vector = op.Reshape(index_vector, shape_vector)
+
+        # Expand reshaped_index_vector to match the full shape of values
+        expanded_index_vector = op.Expand(reshaped_index_vector, values_shape)
+
+        # Flatten into a 1D vector
+        column_index_vector = op.Reshape(expanded_index_vector, [-1])
+
+        # Convert into a column vector to prepare for concatenation
+        column_index_vector = op.Unsqueeze(column_index_vector, axes=[1])
+        index_vectors.append(column_index_vector)
+
+    # Contains all indices to be upadated
+    new_index = op.Concat(*index_vectors, axis=1)
+
+    # Flatten values to match the indices
+    flat_values = op.Reshape(values, [-1])
 
     if accumulate:
-        result = op.ScatterND(self, new_index, values, reduction="add")
+        result = op.ScatterND(self, new_index, flat_values, reduction="add")
     else:
-        result = op.ScatterND(self, new_index, values)
+        result = op.ScatterND(self, new_index, flat_values)
 
     return result
 

tests/function_libs/torch_lib/extra_opinfo.py

@@ -790,20 +790,27 @@ def sample_inputs_index_put(op_info, device, dtype, requires_grad, **kwargs):
     del op_info
     del kwargs
 
-    data = torch_testing.make_tensor(
-        (10, 3),
-        device=device,
-        dtype=dtype,
-        requires_grad=requires_grad,
-    )
-    indices = [torch.arange(8, dtype=torch.int64, device=device).reshape((-1, 4))]
-    values = torch_testing.make_tensor(
-        (2, 4, 3),
-        device=device,
-        dtype=dtype,
-        requires_grad=requires_grad,
+    make_arg = functools.partial(
+        torch_testing.make_tensor, device=device, dtype=dtype, requires_grad=requires_grad
     )
-    yield opinfo_core.SampleInput(data, indices, values)
+
+    data = make_arg((10, 3, 4))
+
+    cases = [
+        # Case 1: Full slices in dims 0 and 2, tensor index in dim 1
+        ([None, torch.arange(2, device=device), None], (10, 2, 4)),
+        # Case 2: Tensor index in dim 0, full slices in dims 1 and 2
+        ([torch.arange(5, device=device), None, None], (5, 3, 4)),
+        # Case 3: Full slices in dims 0 and 1, tensor index in dim 2
+        ([None, None, torch.arange(3, device=device)], (10, 3, 3)),
+        # Case 4: Single index in last dimension
+        ([None, None, torch.tensor([0], device=device)], (10, 3, 1)),
+    ]
+
+    for indices, values_shape in cases:  # type: ignore[misc]
+        values = make_arg(values_shape)  # type: ignore[has-type]
+
+        yield opinfo_core.SampleInput(data, indices, values)
 
 
 def sample_inputs_layer_norm(op_info, device, dtype, requires_grad, **kwargs):