Jitted function: output tensor shape depends on input tensor values

[refraction-ray]

I know well that "output tensor shape depends on input tensor values" paradigm is forbidden in jax jitted function (mentioned everywhere in issues and docs), but still curious on whether there is a workaround (not the static_argnums one, since the output shapes can vary a lot which makes the function retrace each time ).

After all, tensorflow can do this paradigm fine even with jit_compile=True enabled, i.e. compiling to XLA. See the demo below. If there is currently no satisfying workaround in jax, any future plan on such a feature?

@tf.function(jit_compile=True)
def f(a):
    print("tracing")
    return tf.ones(a)

print(f(tf.constant(1)))
print(f(tf.constant(3)))

# tracing # not how tracing is only called once
# tf.Tensor([1.], shape=(1,), dtype=float32)
# tf.Tensor([1. 1. 1.], shape=(3,), dtype=float32)

@jax.jit
def f(a):
    print("tracing")
    return jnp.ones(a)
print(f(jnp.array(1)))
print(f(jnp.array(3)))

# of course, TypeError, the familiar "Shapes must be 1D sequences of concrete values of integer type"

[jakevdp]

Yes, you've read the documentation correctly: jax transforms (including JIT) require arrays to be statically shaped. There is some ongoing experimentation with dynamic shapes, but it is not yet part of the package.