jax\_utils package
==================

Submodules
----------

jax\_utils.common\_tensor module
--------------------------------

.. automodule:: jax_utils.common_tensor
   :members:
   :undoc-members:
   :show-inheritance:

jax\_utils.compilation module
-----------------------------

.. automodule:: jax_utils.compilation
   :members:
   :undoc-members:
   :show-inheritance:

jax\_utils.dynamics module
--------------------------

.. automodule:: jax_utils.dynamics
   :members:
   :undoc-members:
   :show-inheritance:

jax\_utils.gradient module
--------------------------

.. automodule:: jax_utils.gradient
   :members:
   :undoc-members:
   :show-inheritance:

jax\_utils.jax\_tensor module
-----------------------------

.. automodule:: jax_utils.jax_tensor
   :members:
   :undoc-members:
   :show-inheritance:

jax\_utils.markov\_decision\_process module
-------------------------------------------

.. automodule:: jax_utils.markov_decision_process
   :members:
   :undoc-members:
   :show-inheritance:

jax\_utils.optim module
-----------------------

.. automodule:: jax_utils.optim
   :members:
   :undoc-members:
   :show-inheritance:

jax\_utils.pytree module
------------------------

.. automodule:: jax_utils.pytree
   :members:
   :undoc-members:
   :show-inheritance:

jax\_utils.tranform module
--------------------------

.. automodule:: jax_utils.tranform
   :members:
   :undoc-members:
   :show-inheritance:

jax\_utils.typing module
------------------------

.. automodule:: jax_utils.typing
   :members:
   :undoc-members:
   :show-inheritance:

jax\_utils.vectorization module
-------------------------------

.. automodule:: jax_utils.vectorization
   :members:
   :undoc-members:
   :show-inheritance:

Module contents
---------------

.. automodule:: jax_utils
   :members:
   :undoc-members:
   :show-inheritance: