Repository: onnx/onnx
File: onnx/backend/test/case/node/scatterelements.py (Python)

Code before:
    idx = [
        [
            unpack(np.indices(idx_xsection_shape).reshape(indices.ndim - 1, -1)),
            indices[tuple(make_slice(indices, axis, i))].reshape(1, -1)[0],
        ]
        for i in range(indices.shape[axis])
    ]
    idx = list(np.concatenate(idx, axis=1))
    idx.insert(axis, idx.pop())

    # updates_idx is a NumPy advanced indices for indexing of elements in the updates
    updates_idx = list(idx)
    updates_idx.pop(axis)
    updates_idx.insert(
        axis, np.repeat(np.arange(indices.shape[axis]), np.prod(idx_xsection_shape))
    )

    scattered = np.copy(data)
    if reduction == "none":
        scattered[tuple(idx)] = updates[tuple(updates_idx)]
    else:
        idx, updates_idx = (
            make_indices_for_duplicate(idx),
            make_indices_for_duplicate(updates_idx),
        )
        for _iter, idx_set in enumerate(idx):
            if reduction == "add":
                scattered[idx_set] += updates[updates_idx[_iter]]
            elif reduction == "mul":
                scattered[idx_set] *= updates[updates_idx[_iter]]
            elif reduction == "max":
                scattered[idx_set] = np.maximum(
                    scattered[idx_set], updates[updates_idx[_iter]]
                )
            elif reduction == "min":
                scattered[idx_set] = np.minimum(
                    scattered[idx_set], updates[updates_idx[_iter]]
                )
    return scattered


class ScatterElements(Base):
    @staticmethod
    def export_scatter_elements_without_axis() -> None:
        node = onnx.helper.make_node(
            "ScatterElements",
            inputs=["data", "indices", "updates"],
            outputs=["y"],
        )
        data = np.zeros((3, 3), dtype=np.float32)
        indices = np.array([[1, 0, 2], [0, 2, 1]], dtype=np.int64)

Review comment (suggestion):
```suggestion
        for iter_, idx_set in enumerate(idx):
```

Same applies to other changes, could you update? Thank you

Code after:
    idx = [
        [
            unpack(np.indices(idx_xsection_shape).reshape(indices.ndim - 1, -1)),
            indices[tuple(make_slice(indices, axis, i))].reshape(1, -1)[0],
        ]
        for i in range(indices.shape[axis])
    ]
    idx = list(np.concatenate(idx, axis=1))
    idx.insert(axis, idx.pop())

    # updates_idx is a NumPy advanced indices for indexing of elements in the updates
    updates_idx = list(idx)
    updates_idx.pop(axis)
    updates_idx.insert(
        axis, np.repeat(np.arange(indices.shape[axis]), np.prod(idx_xsection_shape))
    )

    scattered = np.copy(data)
    if reduction == "none":
        scattered[tuple(idx)] = updates[tuple(updates_idx)]
    else:
        idx, updates_idx = (
            make_indices_for_duplicate(idx),
            make_indices_for_duplicate(updates_idx),
        )
        for iter_, idx_set in enumerate(idx):
            if reduction == "add":
                scattered[idx_set] += updates[updates_idx[iter_]]
            elif reduction == "mul":
                scattered[idx_set] *= updates[updates_idx[iter_]]
            elif reduction == "max":
                scattered[idx_set] = np.maximum(
                    scattered[idx_set], updates[updates_idx[iter_]]
                )
            elif reduction == "min":
                scattered[idx_set] = np.minimum(
                    scattered[idx_set], updates[updates_idx[iter_]]
                )
    return scattered


class ScatterElements(Base):
    @staticmethod
    def export_scatter_elements_without_axis() -> None:
        node = onnx.helper.make_node(
            "ScatterElements",
            inputs=["data", "indices", "updates"],
            outputs=["y"],
        )
        data = np.zeros((3, 3), dtype=np.float32)
        indices = np.array([[1, 0, 2], [0, 2, 1]], dtype=np.int64)