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.

[unalmis]

As temporary workarounds,

1. Given that JAX does not yet support dynamic shapes, it is typical to operate with a padded array that has junk elements (e.g. zeros). Is it possible to call stop_gradient on those junk elements? That way the shape of the array remains static, but some operations will not enter the Jacobian computational graph, saving memory.
2. My function f which needs dynamic shape support for performance (and hence would benefit from changing to be non jit able currently) is buried inside a function g where g is jitted. When I call g, is there a simple way to execute the function f without jit while doing everything else in g with jit? Perhaps with a pure callback? This would help avoid major rewrite of the code.