Back to jax

``jax.extend.core`` module

docs/jax.extend.core.rst

0.3.25292 B
Original Source

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