What JAX Traces, and What It Refuses

Coming from PyTorch, the first thing JAX does is offend you. You write a simple function with an if statement, wrap it in jax.jit, and call it. Instead of running, it throws a TracerBoolConversionError. No graceful fallback. No partial compile. Just a hard stop on a line of Python that would have run fine a second ago. torch.compile would have shrugged, taken a graph break, and moved on. JAX ...