JitΒΆ

Just-In-Time compilation decorator/function.

Abstract Signature:

Jit(fun: Callable, static_argnums: Optional[Union[int, Sequence[int]]])

PyTorch

API: torch.compile
Strategy: Direct Mapping

JAX (Core)

API: jax.jit
Strategy: Direct Mapping

TensorFlow

API: tf.function
Strategy: Direct Mapping

Apple MLX

API: mlx.core.compile
Strategy: Direct Mapping

Flax NNX

API: flax.nnx.jit
Strategy: Direct Mapping