mirror of
https://github.com/ROCm/jax.git
synced 2025-04-14 10:56: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
16 lines
264 B
ReStructuredText
16 lines
264 B
ReStructuredText
``jax.extend.linear_util`` module
|
|
=================================
|
|
|
|
.. automodule:: jax.extend.linear_util
|
|
|
|
.. autosummary::
|
|
:toctree: _autosummary
|
|
|
|
StoreException
|
|
WrappedFun
|
|
cache
|
|
merge_linear_aux
|
|
transformation
|
|
transformation_with_aux
|
|
wrap_init
|