Content - f737e51ef081e65b285b0ce59d7feb309265f151 - 7e0026a/docs/jax.extend.rst

visit type:

Tip revision: fe0f12e00e53fd9557723c6398797cc4514992f1 authored by Adam Paszke on 18 September 2023, 17:53:10 UTC
[Pallas] Wire up cost estimates in Mosaic params

Tip revision: fe0f12e

jax.extend.rst

``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

Browse the archive

https://github.com/google/jax