``jax.extend`` module
=====================

.. automodule:: jax.extend

``jax.extend.linear_util``
--------------------------

.. automodule:: jax.extend.linear_util

.. autosummary::
  :toctree: _autosummary

  StoreException
  WrappedFun
  cache
  merge_linear_aux
  transformation
  transformation_with_aux
  wrap_init


``jax.extend.random``
---------------------

.. automodule:: jax.extend.random

.. autosummary::
  :toctree: _autosummary

  PRNGImpl
  seed_with_impl
  threefry2x32_p
  threefry_2x32
  threefry_prng_impl
  rbg_prng_impl
  unsafe_rbg_prng_impl
  wrap_key_data