Skip to content

← JAX Performance & Distributed step 10 of 25

Medium Primitives

disable_jit Context Manager

Jump to the task

Run the Python body while debugging

With ordinary JIT, Python executes the function body during tracing. A cached call runs the compiled calculation without executing that Python body again. Python print is not swallowed: it runs during tracing and may show a tracer for a dynamic value.

with jax.disable_jit(): temporarily bypasses JIT behavior for calls made in that context. This lets you use ordinary Python debugging on concrete inputs without removing decorators from your code.

import jax
import jax.numpy as jnp

@jax.jit
def noisy_sum(x):
    print("Python body ran")
    return jnp.sum(x)

x = jnp.array([1., 2., 3.])
noisy_sum(x)  # prints while tracing this signature
noisy_sum(x)  # cached call; no Python print

with jax.disable_jit():
    noisy_sum(x)  # prints
    noisy_sum(x)  # prints again

noisy_sum(x)  # JIT behavior restored; cached call

Every call above returns the scalar value 6. The change is when the Python body runs, not what summing these values means.

What the context does and does not change

The setting applies to nested JIT calls and implicit JIT uses in JAX libraries during the context. On exit, the previous setting is restored. Individual primitive operations can still be compiled by XLA. This is not a switch from JAX arrays to NumPy arrays or from an accelerator to a CPU. Nor does it guarantee a particular thread count or speed relative to NumPy.

It can remove useful compilation optimizations, so restore normal settings before measuring your intended compiled workload. It also does not remove JAX array semantics, such as functional updates, or replace the explicit keys used by jax.random's functional API.

Put the context around the call

When debugging concrete values, enter the context before calling the jitted function. A context inside an outer JIT trace cannot turn its existing tracers into concrete values. Other transformations, including grad and vmap, can also introduce tracers while JIT is disabled.

Python if or assert can inspect a concrete scalar condition in eager execution, but this does not make an arbitrary array a valid boolean. Success with JIT disabled does not establish that the code is compatible with JIT enabled. Recheck the compiled path after diagnosing the issue.

Problem

Implement with_jit_disabled_sum(x) by entering a jax.disable_jit() context and computing jnp.sum(x ** 2) inside it. Return that scalar JAX array.

X is a one-dimensional float32 JAX array. For [3, 4], square to [9, 16] and sum to 25. An empty array returns zero. Do not convert the result to a Python float or change the global configuration permanently.

Call this exercise as an ordinary Python function when demonstrating the context. Wrapping the entire function in an outer JIT would not prove that its arithmetic ran eagerly. Numerical tests verify the sum; inspect the context placement separately.