UndefinedPoly error in jax2tf

[cgarciae]

Hey, I am trying to create a saved model with a variable batch dimension from jax. In my personal code I was getting a UndefinedPoly error with the named batch dimension at some point in the Flax module so I tried running the saved_model_main.py example with --serving_batch_size=-1 just to be sure but I get the same error:

...
File "/data/cristian/jax/jax/experimental/jax2tf/examples/mnist_lib.py", line 108, in predict
x = jnp.dot(x, w) + b
...
File "/data/cristian/elegy/.venv/lib/python3.8/site-packages/jax/interpreters/masking.py", line 224, in eq
raise UndefinedPoly(f"Polynomial comparison {self} == {other} is inconclusive")
jax.interpreters.masking.UndefinedPoly: Polynomial comparison batch == 1 is inconclusive

It seems that Jax is not able to broadcast to "named" dimensions. Creating a saved model with static dimension works just fine.

[gnecula]

Shape polymorphism is not yet supported in jax2tf. We had it working a few months ago but we rolled that support back because there were some unsoundness bugs. We do want to support this again, but it requires some changes in JAX core.

[cgarciae]

Thanks for the info @gnecula ! It might be worth a little note on the examples (unless I missed it). I'll add support only for static dimensions for now.

[gnecula]

I'll keep this issue open, for me to go through the examples and remove mentions of batch polymorphism.

[cgarciae]

Is there an issue for batch polymorphism in general I can follow?

[gnecula]

I created issue #6080