Convert Jax to TF with multiple inputs

[Skyy93]

Hello,

i want to convert a multimodal transformer from jax to TF. The jax2tf convert method does not support two inputs at the same time with the error message: raise ValueError(f"{name} specification must be a tree prefix of the "

tf_fn = jax2tf.convert(
    predict_fn,
    with_gradient=with_gradient,
    polymorphic_shapes=[[None, "b, 224, 224, 3"],[None, "b, 512"]],
)

Did someone encountered this problem and have a solution?

Thank you