Create Ragged Arrays via Vmap

[pharringtonp19]

Related to @KeAWang 's previous post, I am trying to create ragged arrays.

I guess my question in whether this will be possible to do in the future?! -- Thanks

import jax
import jax.numpy as jnp
from jax.interpreters import batching
from functools import partial 

jax.config.update("jax_array", False)
jax.config.update("jax_dynamic_shapes", True)
jax.config.update("jax_numpy_rank_promotion", "allow")

def cluster_data(key, n): 
    return jax.random.normal(key, shape=(n,1))

sizes = jnp.array([3, 1, 4])
zs = jax.vmap(cluster_data, out_axes=batching.pile_axis)(jax.random.split(jax.random.PRNGKey(0), len(sizes)), sizes)
print(zs)

ValueError: Non-hashable static arguments are not supported, as this can lead to unexpected cache-misses. Static argument (index 1) of type <class 'jax._src.core.NamedShape'> for function _normal is non-hashable.

[jakevdp]

@mattjj and @apaszke have done some experimentst with this and may be able to give some details, but there's not yet any general support for ragged arrays in JAX.

[mattjj]

That's right. And many places in the codebase are written assuming static shapes (like in this case, jax.random.normal's use of static_argnums for its shape parameter), so those need to be unwound. Each one is easy but there are a lot of places to cover.

[pharringtonp19]

@jakevdp @mattjj Sounds good! Thanks