JITable indices matching an array's shape

[sharlaon]

We are given arr = jax.numpy.array(...) where a.shape can be anything, and we seek another array idxs whose shape is (*a.shape, 1) having roughly the property that idxs[i] == i for all valid arr-indices i. The precise intended behavior is that of the following code:

def idxs(shape):
    return jnp.transpose(
        jnp.indices(shape),
        axes=jnp.roll(jnp.arange(len(shape) + 1), -1)
    )

Thus by including idxs(arr.shape) along with arr as arguments to a vmap, it is useful anywhere one wants to know the index at which they are operating, like having enumerate in standard Python.

The problem with the above idxs code is that, although its shape behavior is obviously statically determined by that of its argument arr, the construction passes through arr.shape, which trips up the JIT compiler. There are functions zeros_like, full_like, and so on, that directly take arr (not arr.shape) and are respected by the JIT compiler. Is there some analogue ("idxs_like") here?

[jakevdp]

which trips up the JIT compiler.

Can you elaborate here? How does this trip up the compiler? Note that zeros_like(x) and friends are basically just a wrapper around lax.full(x.shape, ...) so I'm not sure how that differs from the shape-based function you're defining.

[jakevdp]

Oh, I see from your code example that maybe the issue is you're passing traced arrays to the axes argument. It should work if you do your shape manipulations on host with NumPy rather than on device with JAX:

from functools import partial
import jax
import jax.numpy as jnp
import numpy as np

@partial(jax.jit, static_argnames=['shape'])
def idxs(shape):
    return jnp.transpose(
        jnp.indices(shape),
        axes=np.roll(np.arange(len(shape) + 1), -1)
    )
print(idxs((4,)))

[[0]
 [1]
 [2]
 [3]]