Why TrainState can be use as arguments of jited function without trigger error?

[wztdream]

Hi,
Just for curiosity, as jit doc says the arguments of a jit-ed function must be:

Its arguments and return value should be arrays,
scalars, or (nested) standard Python containers (tuple/list/dict) thereof

But TrainState can be use in jit-ed function without trigger any error, for example:

@functools.partial(jax.jit, static_argnums=(2, 3))
def train_step(state, batch, loss_fn, train_metrics_fn, rngs):
    # some train step code here,state is an instance of TrainState

We have to mark loss_fn which is a Callable as static arguments to avoid error, but no need forstate which is also violate the constrains in jit doc, so why there is no error for state?

[marcvanzee]

When you look at the definition of the class TrainState, you can see that it extends struct.PyTreeNode. This means that it is a dataclasses that acts act like a JAX pytree node, so JAX knows how to flatten/unflatten it for JAX transformations.