Open tchatow opened 1 month ago
Thanks for reporting this! Any chance you can share a repro?
Here's a simple example
N, = jax.export.symbolic_shape("N")
f = jax.export.export(jax.jit(jnp.digitize))
shape0 = jax.ShapeDtypeStruct((10,), jnp.int32)
shape1 = jax.ShapeDtypeStruct((N,), jnp.int32)
f(shape0, shape1)
This is trickier than I thought; if we fix that local issue (e.g. by replacing len(a)
with a.shape[0]
) then we get a downstream issue that the searchsorted implementations rely on static sizes (or at least size bounds).
@jakevdp any ideas?
Not sure... it's still not clear to me to what extent we should expect shape polymorphism to be supported in JAX APIs. Do we have those goals documented anywhere? It's pretty incomplete at the moment: if we opened issues like this one for every numpy API that doesn't support shape polymophism, we wouldn't have time to do anything else 😀
Description
Error when tracing with shape polymorphism in jax.numpy.digitize
System info (python version, jaxlib version, accelerator, etc.)