flax

1d669bc5 - [JAX] Fix up uses of PyTree and PyTreeDef types.

[JAX] Fix up uses of PyTree and PyTreeDef types. In many cases, PyTreeDef was used as a type when pytree (Any) should have been used. PiperOrigin-RevId: 515132058