Add arbitrary amount of axis

[pablo2909]

Hi,

I am wondering if there is a way of adding an arbitrary amount of axis to an array X based on an array Z? I would like to do the following in an einops way:

import jax.numpy as jnp

X = jnp.arange(10)
Z = jnp.arange(10*3*4).reshape(10,3,4)

X.reshape([-1] + [1]*(Z.ndim-1))

Thank you for any suggestions

[arogozhnikov]

Hi Paul, einops does not provide such tool, you can use some trick like:

result = (X.T + Z.T).T

[pablo2909]

Thank you ! Do you have an idea why this works though ? and not just X + Y.

[arogozhnikov]

broadcasting can automatically add arbitrary number of dummy axes on the left (but not on the right)

[pablo2909]

did not know that, thank you ! :)