timeline for differentiable code with dynamic shapes?

[unalmis]

Is there a timeline for when JAX will support Dynamic shapes and Jagged tensor operations?

In the near term, is there a way to replace this logic

# this computes the expensive function for each element in x
y = jnp.where(x != 0, expensive_function(x), dummy) 

with this

# this applies expensive function on nonzero elements only
y = nonzero_where(x, expensive_function, dummy)

in a JAX transformable manner? The simple implementations for this in #1303 are not ideal.