fix(index_put): handle None cases

Author: AyoubMDL

onnxscript/function_libs/torch_lib/ops/core.py

@@ -4259,15 +4259,48 @@ 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))
 
-    # TODO(justinchuby): Handle when indicies has more than one element
-    index = indices[0]
-    new_index = op.Unsqueeze(index, [-1])
+    values_shape = values.shape.numpy()
+
+    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
+
+        # 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