vmap with ragged data (list of variable length arrays)

[8bitmp3]

Bumped into a question by @iskandr (I think)

how do use vmap with ragged data (list of variable length arrays)?

Thought I'd share it here and try to get an answer to build up the "knowledge base" under github/jax 🤔

Source: https://twitter.com/iskander/status/1296906143332261889

[charlessievers]

Not an answer, but I would like to know if there are any ways to vmap and handle lists of variable length arrays.

[jakevdp]

JAX does not currently have any facility for computation on ragged arrays, and the vmap transform in particular is built with homogeneous batches in mind. Depending on your use case, you may be able to do what you have in mind using JAX's experimental sparse array support; but although similar in spirit, sparse arrays are not precisely equivalent to ragged arrays.

[charlessievers]

Thanks for the speedy reply! I will give it a look.

[cevheck]

@jakevdp,
Any updates on this issue? I would like to perform a certain calculation on different machine learning models (ensemble method) and do them together in a single vmap call. My issue is that each of those models has a different length of parameters (e.g. combining a single layer 16neurons network with a double layer 32 neurons network).

Thanks in advance!

[jakevdp]

There's been more work on ragged arrays recently (e.g. #16072 and #16324), but it's still all very experimental. @axch and @mattjj know more.