Regression of JAX duck typing.

[AlexeyKurakin]

When JAX variables are used without .value with jit, gradient and other transformation it leads to error TypeError: Value Traced<...> with type <class 'jax.interpreters.partial_eval.DynamicJaxprTracer'> is not a valid JAX type
This is a recent regression caused by some underlying changes in JAX

[AlexeyKurakin]

Temporary workaround - either use .value or rollback to JAX version 0.2.10

[yetiansh]

Hi, I have also encountered this issue when running HuggingFace T5 examples with Jax and Flax. Is there any updates?