Possibly a bug in jax.numpy.ravel_multi_index

[PgLoLo]

Somewhy I can not vectorize jax.numpy.ravel_multi_index with a fixed dims argument. Consider the following code:

indices = jnp.array([[0, 0, 0], [1, 1, 1], [2, 2, 2]])

my_ravel = lambda x: jnp.ravel_multi_index(x, (3, 3, 3))

print(my_ravel(indices[1]))
print(vmap(my_ravel)(indices))

This code throws ConcretizationTypeError on the second print, although there are no arguments to vmapped ravel that should produce such a result. Consider, for example, the following implementation of multi_index ravel:

def ravel_by_hands(multi_index, shapes):
    raveled_index = 0
    for dim, index in zip(shapes, multi_index):
        raveled_index = raveled_index * dim + index
    return raveled_index


indices = jnp.array([[0, 0, 0], [1, 1, 1], [2, 2, 2]])

my_ravel = lambda x: ravel_by_hands(x, (3, 3, 3))

print(my_ravel(indices[1]))
print(vmap(my_ravel)(indices))

This works perfectly fine.

Is there a bug in jax.numpy.ravel_multi_index function or do I misunderstand something?

[jakevdp]

The error looks like this:

ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(int32[])>with<BatchTrace(level=1/0)> with
  val = Array([0, 1, 2], dtype=int32)
  batch_dim = 0
The error occurred because ravel_multi_index was jit-compiled with mode='raise'. Use mode='wrap' or mode='clip' instead.
This BatchTracer with object id 140570390572544 was created on line:
  <ipython-input-2-5c78fe181e6a>:6 (<lambda>)

See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.ConcretizationTypeError

The issue is that the default mode='raise' is not compatible with JAX transforms like jit or vmap, and as mentioned by the error you should use mode='wrap' or mode='clip' instead, whichever is appropriate to your situation:

my_ravel = lambda x: jnp.ravel_multi_index(x, (3, 3, 3), mode='clip')

print(my_ravel(indices[1]))
# 13
print(vmap(my_ravel)(indices))
# [ 0 13 26]

You might ask why did we make the choice to use a default parameter value that is incompatible with JAX transforms: it's because this is a jax wrapper of a numpy API whose default semantics are incompatible with JAX transforms, so our options are either (1) change the default semantics (which may be confusing) or (2) raise an informative error message if the default case is transformed. We opted for 2. Does that make sense?