mirror of
https://github.com/ROCm/jax.git
synced 2025-04-15 19:36:06 +00:00
78543f7bb8
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
17 lines
213 B
ReStructuredText
17 lines
213 B
ReStructuredText
.. currentmodule:: jax.extend
|
|
|
|
``jax.extend`` module
|
|
=====================
|
|
|
|
.. automodule:: jax.extend
|
|
|
|
Modules
|
|
-------
|
|
|
|
.. toctree::
|
|
:maxdepth: 1
|
|
|
|
jax.extend.linear_util
|
|
jax.extend.mlir
|
|
jax.extend.random
|