jax2tf polymorphic_shapes with structured input

I’m trying to use polymorphic shapes in jax2tf. From the docs it seems that structured inputs are supported:

It should be a Python object with the same pytree structure as, or a prefix of, the tuple of arguments to the function, but with a shape specification corresponding to each argument.

yet when I was trying to use it with a NamedTuple I got the error: InconclusiveDimensionOperation: Dimension polynomial 'b' is not constant

Did I do something wrong? Is this supported?

from jax.experimental import jax2tf
import tensorflow.compat.v1 as tf
from typing import Tuple, NamedTuple

class G(NamedTuple):
  a: Tuple[int]
  b: Tuple[int]

def fn(a):
  return a.a[0] + a.b[0]

g = jax2tf.convert(fn, polymorphic_shapes=["(b, ...)"])
h = jax2tf.convert(fn, polymorphic_shapes=[G(a="(b, ...)", b="(b, ...)")])
l = jax2tf.convert(fn, polymorphic_shapes=[G(a="(b, _)", b="(b, _)")])
m = jax2tf.convert(fn, polymorphic_shapes=["(b, _)"])
place = G(a=tf.placeholder(tf.int32, shape=(10, 5)), b=tf.placeholder(tf.int32, shape=(10, 5)))

g(place), h(place), ...

[gnecula]

I was able to reproduce your use case, and indeed PR #7317 should solve this issue.

I think that you should be able to patch this PR locally if you clone the JAX sources, then patch the PR and then run pip install -e . in the JAX directory.

[gnecula]

I asked “why do you think your code is shape polymorphic” because I did not not remember that your function is the result of jax.vmap. In that case we aim to extend JAX to recognise that shape polymorphism. It is possible that we are not there yet, as you seem to be encountering.

From the stack trace, it seems that the problem is in the invocation of lax.dynamic_slice (line 888). Can you get the shapes of the arguments of lax.dynamic_slice? I can try in parallel to construct a repro example.