Assign to diagonal

[AllanYangZhou]

Suppose X is a rank-5 tensor and y a matrix. The operation I want to accomplish is:

X[i,i,i,j,k] += y[j,k]

for all values of indices i,j,k. That is, I want to broadcast the matrix y along the "diagonal" of the first 3 dimensions of X.

More generally, I want to assign arbitrary-rank y to the diagonal along an arbitrary subset of dimensions of X, as long as the shapes are all compatible. Is this possible with any einops-like functions? I could not figure out any way.

Is there some small modification of existing einops implementations that could accomplish this?

[arogozhnikov]

In eindex you can do something like this:

from eindex.numpy import scatter

i, j = x.shape
result = scatter(
    x, [np.arange(i)] * 2, 
    'j k, [i1,i2] i -> i i1 i2 j k',
    i1=i, i2=i,
)

But eindex supports only numpy, as torch/tf/jax still do not support array api.

[AllanYangZhou]

Thanks, eindex looks really promising! Though yeah the torch/jax incompatibility for now might be a bit limiting.