All modules for which code is available

• jax_utils.common_tensor
• jax_utils.compilation
• jax_utils.dynamics
• jax_utils.gradient
• jax_utils.jax_tensor
• jax_utils.markov_decision_process
• jax_utils.optim
• jax_utils.pytree
• jax_utils.tranform
• jax_utils.typing
• jax_utils.vectorization
• jaxlib.xla_extension