rocm_jax/docs/jax.extend.rst

.. currentmodule:: jax.extend

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

.. automodule:: jax.extend

Modules
-------

.. toctree::
  :maxdepth: 1

  jax.extend.linear_util
  jax.extend.mlir
  jax.extend.random
Add jax.extend.mlir. Some users of JAX want to use the MLIR dialects defined in jaxlib. In particular, these need to be used by custom lowering rules. Add a semi-public (jax.extend) API to access these, rather than having them use jax._src.lib.mlir. PiperOrigin-RevId: 588448489 2023-12-06 09:16:04 -08:00			`.. currentmodule:: jax.extend`

initiate `jax.extend` via docs and top-level module set-up 2023-05-15 13:52:28 -07:00			``jax.extend`` module
			`=====================`

			`.. automodule:: jax.extend`
DOC: add missing docs for jax.random functions 2023-08-28 13:55:33 -07:00
Add jax.extend.mlir. Some users of JAX want to use the MLIR dialects defined in jaxlib. In particular, these need to be used by custom lowering rules. Add a semi-public (jax.extend) API to access these, rather than having them use jax._src.lib.mlir. PiperOrigin-RevId: 588448489 2023-12-06 09:16:04 -08:00			`Modules`
			`-------`
DOC: add missing docs for jax.random functions 2023-08-28 13:55:33 -07:00
Add jax.extend.mlir. Some users of JAX want to use the MLIR dialects defined in jaxlib. In particular, these need to be used by custom lowering rules. Add a semi-public (jax.extend) API to access these, rather than having them use jax._src.lib.mlir. PiperOrigin-RevId: 588448489 2023-12-06 09:16:04 -08:00			`.. toctree::`
			`:maxdepth: 1`
DOC: add missing docs for jax.random functions 2023-08-28 13:55:33 -07:00
Add jax.extend.mlir. Some users of JAX want to use the MLIR dialects defined in jaxlib. In particular, these need to be used by custom lowering rules. Add a semi-public (jax.extend) API to access these, rather than having them use jax._src.lib.mlir. PiperOrigin-RevId: 588448489 2023-12-06 09:16:04 -08:00			`jax.extend.linear_util`
			`jax.extend.mlir`
			`jax.extend.random`