jax.Inline

class jax.Inline(value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None)

Enumeration specifying inlining behavior for nested jitted functions.

Values:

JAX_EARLY: Inlined during JAX tracing into enclosing JAXPRs. JAX_LATE: Inlined during during JAX to MLIR lowering. XLA_EARLY: Preserved in MLIR and marked for early inlining by the XLA compiler. XLA_LATE: Preserved in MLIR and marked for late inlining by the XLA compiler. AUTO: Automatic inlining policy determined by JAX and XLA.

__init__(*args, **kwds)

Attributes

JAX_EARLY
JAX_LATE
XLA_EARLY
XLA_LATE
AUTO