Fix scatter for list indices (#982)

**Context:**
Scatter operation is failing lowering when indices are list/tuple.

**Description of the Change:**
The dimensions int vector is not correctly set, this is updated. Also
the inner index for scatter was incorrect, it is corrected as well.

**Benefits:**
More scatter op supported.

doc/changelog.md

@@ -204,6 +204,24 @@
 
 <h3>Bug fixes</h3>
 
+* Fix a bug where scatter did not work correctly with list indices.
+  [(#982)](https://github.com/PennyLaneAI/catalyst/pull/982)
+
+  ```python
+  A = jnp.ones([3, 3]) * 2
+
+  def update(A):
+      A = A.at[[0, 1], :].set(jnp.ones([2, 3]), indices_are_sorted=True, unique_indices=True)
+      return A
+  ```
+
+  ```pycon
+  >>> update
+  [[1. 1. 1.]
+   [1. 1. 1.]
+   [2. 2. 2.]]
+  ```
+
 * Static arguments can now be passed through a QNode when specified
   with the `static_argnums` keyword argument.
   [(#932)](https://github.com/PennyLaneAI/catalyst/pull/932)

frontend/test/pytest/test_scatter.py

@@ -59,6 +59,32 @@ def multiple_index_multiply(l: jax.core.ShapedArray((3, 3), dtype=float)):
     assert np.allclose(res, jnp.array([[0, 1, 2], [9, 12, 15], [18, 21, 24]]))
 
 
+def test_matrix_list_index_right():
+    """Test to index (list) a jax array and have operations on it."""
+
+    A = jnp.ones([3, 3]) * 2
+
+    def matrix_list_index(A):
+        A = A.at[[0, 1], :].set(jnp.ones([2, 3]), indices_are_sorted=True, unique_indices=True)
+        return A
+
+    res = qjit(matrix_list_index)(A)
+    assert np.allclose(res, jnp.array([[1, 1, 1], [1, 1, 1], [2, 2, 2]]))
+
+
+def test_matrix_list_index_left():
+    """Test to index (list) a jax array and have operations on it."""
+
+    A = jnp.ones([3, 3]) * 2
+
+    def matrix_list_index(A):
+        A = A.at[:, [0, 1]].set(jnp.ones([3, 2]), indices_are_sorted=True, unique_indices=True)
+        return A
+
+    res = qjit(matrix_list_index)(A)
+    assert np.allclose(res, jnp.array([[1, 1, 2], [1, 1, 2], [1, 1, 2]]))
+
+
 def test_gather_derivative():
     """Test the derivative of indexing."""
 

mlir/lib/Catalyst/Transforms/ScatterPatterns.cpp

@@ -228,16 +228,13 @@ struct ScatterOpRewritePattern : public mlir::OpRewritePattern<mhlo::ScatterOp>
         int64_t start = 0;
         std::vector<int64_t> dimensions;
 
-        for (int64_t i = start; i <= updatesSize; ++i) {
+        for (int64_t i = start; i < updatesSize; ++i) {
             dimensions.push_back(i);
         }
         if (!data.updatedWindowsDims.empty()) {
-            std::copy_if(data.updatedWindowsDims.begin(), data.updatedWindowsDims.end(),
-                         std::back_inserter(data.updatedScatterDims),
-                         [&dimensions](int64_t element) {
-                             return std::find(dimensions.begin(), dimensions.end(), element) !=
-                                    dimensions.end();
-                         });
+            std::set_difference(dimensions.begin(), dimensions.end(),
+                                data.updatedWindowsDims.begin(), data.updatedWindowsDims.end(),
+                                std::back_inserter(data.updatedScatterDims));
         }
         else {
             data.updatedScatterDims = dimensions;
@@ -364,9 +361,10 @@ struct ScatterOpRewritePattern : public mlir::OpRewritePattern<mhlo::ScatterOp>
                 auto itScatter =
                     std::find(scatterDimsToOperandDims.begin(), scatterDimsToOperandDims.end(), i);
                 if (itScatter != scatterDimsToOperandDims.end()) {
-                    Value index = builder.create<index::ConstantOp>(loc, i);
+                    int innerIndex = std::distance(scatterDimsToOperandDims.begin(), itScatter);
+                    Value indexConstantOp = builder.create<index::ConstantOp>(loc, innerIndex);
                     auto indexScatter =
-                        builder.create<tensor::ExtractOp>(loc, scatterIndices, index);
+                        builder.create<tensor::ExtractOp>(loc, scatterIndices, indexConstantOp);
                     auto indexUpdateCasted =
                         builder.create<index::CastSOp>(loc, indexScatter.getType(), indexUpdate);
                     Value addValue =
@@ -388,9 +386,10 @@ struct ScatterOpRewritePattern : public mlir::OpRewritePattern<mhlo::ScatterOp>
                 auto itScatter =
                     std::find(scatterDimsToOperandDims.begin(), scatterDimsToOperandDims.end(), i);
                 if (itScatter != scatterDimsToOperandDims.end()) {
-                    Value index = builder.create<index::ConstantOp>(loc, i);
+                    int innerIndex = std::distance(scatterDimsToOperandDims.begin(), itScatter);
+                    Value indexConstantOp = builder.create<index::ConstantOp>(loc, innerIndex);
                     auto indexScatter =
-                        builder.create<tensor::ExtractOp>(loc, scatterIndices, index);
+                        builder.create<tensor::ExtractOp>(loc, scatterIndices, indexConstantOp);
                     fullStartIndex.push_back(indexScatter);
                 }
                 else {
@@ -400,16 +399,16 @@ struct ScatterOpRewritePattern : public mlir::OpRewritePattern<mhlo::ScatterOp>
                 }
             }
             // Full windows indices
-            SmallVector<Value> fullWindowIndex;
+            SmallVector<Value> fullWindowIndex = updateWindowsIndices;
             for (auto insertedDim : insertedWindowsDims) {
                 auto c0 = builder.create<index::ConstantOp>(loc, 0);
-                updateWindowsIndices.insert(updateWindowsIndices.begin() + insertedDim, c0);
+                fullWindowIndex.insert(fullWindowIndex.begin() + insertedDim, c0);
             }
             // Add
             SmallVector<Value> results;
-            for (size_t i = 0; i < updateWindowsIndices.size(); ++i) {
+            for (size_t i = 0; i < fullWindowIndex.size(); ++i) {
                 Value indexScatter = fullStartIndex[i];
-                auto indexUpdate = updateWindowsIndices[i];
+                Value indexUpdate = fullWindowIndex[i];
                 auto indexUpdateCasted =
                     builder.create<index::CastSOp>(loc, indexScatter.getType(), indexUpdate);
                 Value addValue =