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 =