jax.extend.core module

.. automodule:: jax.extend.core

.. autosummary:: :toctree: _autosummary

ClosedJaxpr Jaxpr JaxprEqn Literal Primitive Token Var array_types jaxpr_as_fun primitives mapped_aval take_current_trace unmapped_aval