Full coverage for __jax_array__ protocol

[maxencefaldor]

Description

I am using Flax with the NNX API that requires arrays to be encapsulated within nnx.Param.

The issue is that NNX is using the __jax_array__ protocol to avoid users having to type .value but JAX doesn't have full coverage for __jax_array__.

For example, this code from an NNX model raises the following error:

self.reshape_c_k = nnx.Param(reshape_c_k)
state_fft_k = jnp.dot(state_fft, self.reshape_c_k)

Error:

TypeError: Argument 'Param(value=Traced<ShapedArray(float32[3,3])>with<DynamicJaxprTrace(level=3/0)>)'
of type <class 'flax.nnx.variablelib.Param'> is not a valid JAX type.

Maybe @cgarciae can provide more insights on this if necessary?

System info (python version, jaxlib version, accelerator, etc.)

python version: 3.12.6
jax: 0.4.34
jaxlib: 0.4.34

[jakevdp]

Hi - thanks for the request. This issue has been discussed previously in #10065. But the short version is that we never meant __jax_array__ to be a fully-supported public API, and it's not clear that it's the best mechanism for downstream libraries to rely on. cc/ @mattjj who has more context.