diff --git a/doc/changelog.md b/doc/changelog.md
index eecd735ed4..7dc0b1a07a 100644
--- a/doc/changelog.md
+++ b/doc/changelog.md
@@ -204,6 +204,24 @@
Bug fixes
+* 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)
diff --git a/frontend/test/pytest/test_scatter.py b/frontend/test/pytest/test_scatter.py
index 1380079e31..06b3cd5075 100644
--- a/frontend/test/pytest/test_scatter.py
+++ b/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."""
diff --git a/mlir/lib/Catalyst/Transforms/ScatterPatterns.cpp b/mlir/lib/Catalyst/Transforms/ScatterPatterns.cpp
index c4013e5ee7..013cdec2a3 100644
--- a/mlir/lib/Catalyst/Transforms/ScatterPatterns.cpp
+++ b/mlir/lib/Catalyst/Transforms/ScatterPatterns.cpp
@@ -228,16 +228,13 @@ struct ScatterOpRewritePattern : public mlir::OpRewritePattern
int64_t start = 0;
std::vector 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
auto itScatter =
std::find(scatterDimsToOperandDims.begin(), scatterDimsToOperandDims.end(), i);
if (itScatter != scatterDimsToOperandDims.end()) {
- Value index = builder.create(loc, i);
+ int innerIndex = std::distance(scatterDimsToOperandDims.begin(), itScatter);
+ Value indexConstantOp = builder.create(loc, innerIndex);
auto indexScatter =
- builder.create(loc, scatterIndices, index);
+ builder.create(loc, scatterIndices, indexConstantOp);
auto indexUpdateCasted =
builder.create(loc, indexScatter.getType(), indexUpdate);
Value addValue =
@@ -388,9 +386,10 @@ struct ScatterOpRewritePattern : public mlir::OpRewritePattern
auto itScatter =
std::find(scatterDimsToOperandDims.begin(), scatterDimsToOperandDims.end(), i);
if (itScatter != scatterDimsToOperandDims.end()) {
- Value index = builder.create(loc, i);
+ int innerIndex = std::distance(scatterDimsToOperandDims.begin(), itScatter);
+ Value indexConstantOp = builder.create(loc, innerIndex);
auto indexScatter =
- builder.create(loc, scatterIndices, index);
+ builder.create(loc, scatterIndices, indexConstantOp);
fullStartIndex.push_back(indexScatter);
}
else {
@@ -400,16 +399,16 @@ struct ScatterOpRewritePattern : public mlir::OpRewritePattern
}
}
// Full windows indices
- SmallVector fullWindowIndex;
+ SmallVector fullWindowIndex = updateWindowsIndices;
for (auto insertedDim : insertedWindowsDims) {
auto c0 = builder.create(loc, 0);
- updateWindowsIndices.insert(updateWindowsIndices.begin() + insertedDim, c0);
+ fullWindowIndex.insert(fullWindowIndex.begin() + insertedDim, c0);
}
// Add
SmallVector 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(loc, indexScatter.getType(), indexUpdate);
Value addValue =