jax.flatten_util module .. currentmodule:: jax.flatten_util .. automodule:: jax.flatten_util List of Functions .. autosummary:: :toctree: _autosummary ravel_pytree