mlx equivalent to np.add.at / torch.Tensor.scatter_add_

[louisamand]

Would it be complicated to make an equivalent of np.add.at and torch.Tensor.scatter_add_ (as shown in the example bellow) in mlx ?

import numpy as np
import torch

idx = np.array([0, 1, 2, 2, 1, 0])
b = np.array([0.1, 9, -1, 4, 3, 2])

# loop
a = np.zeros(shape=3)
for i, k in zip(idx, b):
    a[i] += k
print(a) # [ 2.1 12.   3. ]

# numpy
a = np.zeros(shape=3)
np.add.at(a, idx, b)
print(a) # [ 2.1 12.   3. ]

# torch
idx = torch.from_numpy(idx)
b = torch.from_numpy(b)
a = torch.zeros(3).double().scatter_add_(dim=0, index=idx, src=b)
print(a) # tensor([ 2.1000, 12.0000,  3.0000], dtype=torch.float64)

[awni]

We have something like a scatter add in MLX:

import mlx.core as mx

idx = mx.array([0, 1, 2, 2, 1, 0])
b = mx.array([0.1, 9, -1, 4, 3, 2])
a = mx.zeros(shape=3)

a = a.at[idx].add(b)

docs: https://ml-explore.github.io/mlx/build/html/python/_autosummary/mlx.core.array.at.html#mlx.core.array.at

[louisamand]

Thanks a lot!