Status of the JAX dynamic API

[sergei-mironov]

Hello!
Could you please share your status/plans regarding the JAX dynamic API, covered by this test? I am particularly interested in array indexing operations.

To give some context, I am exploring options for the dynamic array support in the Catalyst quantum compiler where we use jax as a frontend and a modified StableHLO lowering pipeline as a backend. I checked the JAX 0.4.19 against some common indexing use-cases, see the table below. Can you confirm the correctness of the approach in general? Should we expect improvements in this area in the near future?

Use-case	Got JAXPR? [1]	Got MLIR? [2]	Compiles and runs? [3]
Indexing with constant arr[42]	yes	yes	yes
Indexing with variable arr[idx]	yes	yes	yes
Indexing with colon in one dim arr[0,:,0]	yes	no	-
Indexing with constant slice arr[0:42]	no	-	-
Indexing with variable slice arr[0:idx]	no	-	-
Partial indexing jnp.ones([4,4])[0]	yes	no	-

• [1,2] - The detailed output is here. Note that I had to make a few hacks in JAX in order to get JAXPRs. See the patch linked below.
• [3] - Our custom MLIR lowering pipeline lowers StableHLO to LLVM and then to x86 and runs it using
  a custom runtime

All the use-cases I tried are of the following form:

import jax
from catalyst.utils import jaxpr_to_mlir, compile_and_run_mlir # our custom utilities

jax.config.update("jax_dynamic_shapes", True)

def func(sz:int):
  arr = jnp.ones((sz,sz,sz), jnp.float32)
  return arr[0,:,0]

jaxpr = jax.make_jaxpr(func)(3)
print(jaxpr)
mlir = jaxpr_to_mlir("test", jaxpr, jaxpr.out_avals[0].shape)
print(mlir)
res = compile_and_run_mlir(mlir)(3)
print(res)

Where:

• jaxpr_to_mlir is our thin wrapper around lower_jaxpr_to_fun
  I think it is still representative.
• compile_and_run_mlir runs our customized StableHLO lowering to LLVM. It might contain unrelated
  issues, I am showing it here for reference.
• JAX version is 0.4.19, commit 604dc24 with a few hacks

Updates

• Removed the incorrect result notice from point [3], the issue was on our side
• Removed the arr[:] case from the table due to an error in test case

[jakevdp]

Hi - I'm a bit confused by what you're asking here. One clarification: it looks like most of your examples fail because they try to create a dynamically-sized array before slicing/indexing it. Is that intentional?

For example, this fails:

def func(sz:int):
  o = jnp.ones(sz, jnp.float32)
  return o[0:2]
jaxpr = jax.make_jaxpr(func)(3)

But that failure has nothing to do with indexing; if you perform the indexing on a statically-sized array, it's fine:

def func(o):
  return o[0:2]
o = jnp.ones(3.0)
jaxpr = jax.make_jaxpr(func)(o)

[sergei-mironov]

@jakevdp

it looks like most of your examples fail because they try to create a dynamically-sized array before slicing/indexing it. Is that intentional?

Yes, indeed! I am interested to know the status of the dynamic API in general, and the situation around indexing of dynamically-shaped arrays in particular. Static indexing is out of the scope.

[jakevdp]

Oh, I see, thanks – the dynamic API is experimental and incomplete, and likely to remain that way for the foreseeable future. The things that work work, and the things that don't don't, and you should expect it to remain in that state.

[daskol]

@jakevdp What is the main obstacles in completing dynamic shapes support? It would be great to see a road map if there is any. There is some progress in last years but it is unclear should we wait a little bit more or it is better to switch to other platform like Julia as a particular example.

[jakevdp]

@mattjj might be able to comment on that