{
  "source_file": "0930",
  "source_platform": "Hashnode",
  "author": "Wesley Kambale",
  "title_listed": "When JIT hides your errors: Debugging JAX",
  "keywords": "",
  "domain": "机器学习",
  "content_type": "深度技术解析",
  "is_paid_flag": "否",
  "url": "https://kambale.dev/debugging-jax",
  "publish_date": "2026-07-18",
  "corpus_stage": "",
  "notes": "",
  "id": "487b29700e83",
  "collected_at": "2026-10-02T14:10:59",
  "title": "When JIT hides your errors: Debugging JAX",
  "body_markdown": "When JIT hides your errors: Debugging JAX\nUpdated\nJuly 29, 2026\n•\n13\nmin read\n•\nView as Markdown\nW\nWesley Kambale\nI'm a Machine Learning Engineer passionate about building production-ready ML systems for the African market.\nWith experience in TensorFlow, Keras, and Python-based workflows, I help teams bridge the gap between machine learning research and real-world deployment—especially on resource-constrained devices.\nI'm also a Google Developer Expert in AI. I regularly speak at tech conferences including PyCon Africa, DevFest Kampala, DevFest Nairobi and more and also write technical articles on AI/ML here.\nPart of series\nJAX\nLast week,\nChex\ngave us a firewall against shape mismatches and NaN gradients. That firewall catches the bugs that announce themselves; a tensor with the wrong shape, a gradient that explodes into infinity.\nBut not every bug is loud. Sometimes your model trains without a single error. The loss goes down. The shapes all line up. And yet the predictions are subtly, mysteriously wrong. Somewhere in your forward pass, the logic itself is flawed; and there's no assertion for \"this math is conceptually incorrect.\"\nWhen that happens, your first instinct is probably the one every Python developer reaches for: sprinkle in a few\nprint()\nstatements, or drop a\nbreakpoint()\nin the middle of the function. In standard Python, this works instantly.\nInside a\n@jax.jit\n-compiled function, it doesn't work the way you expect, and understanding exactly why is the key to debugging JAX effectively.\nToday we open up the compiled black box. By the end of this article, you'll know how to:\nUnderstand why standard debugging tools fail inside JIT\nPrint real runtime values from inside compiled functions with\njax.debug.print\nPause execution mid-computation with\njax.debug.breakpoint\nTemporarily disable compilation entirely to use\npdb\nlike normal Python\nAutomatically catch the exact operation that produces a NaN\nUse\nnnx.display()\nand\nModule.sow()\nto inspect Flax NNX models\nCombine all of these into a systematic debugging workflow\nLet's go hunting.\nWhy your debugger doesn't work inside JIT\nTo understand JAX debugging, you have to remember the two-phase life of a\n@jax.jit\nfunction, which we first touched on in\nweek 1\n.\nPhase 1: Tracing.\nThe first time you call a JIT-compiled function, JAX runs it once; not with your real data, but with abstract\ntracers\n. A tracer knows the shape and dtype of an eventual array, but not its actual values. JAX uses this trace to build a computation graph.\nPhase 2: Compilation and Execution.\nXLA takes that graph, compiles it into optimized machine code, and only\nthen\ndoes your real data flow through it; as a single fused kernel, without Python in the loop at all.\nHere's the problem for debugging: standard Python tools assume every line of code executes with real, inspectable values, one at a time, with Python in control the whole way through.\ndef broken_debug(x):\n    print(f\"x = {x}\")       # Prints an abstract tracer, not a value\n    if x > 0:                # Can crash, tracers don't have a boolean truth value\n        return x * 2\n    return x\n\njax.jit(broken_debug)(5.0)\nRun this, and two things go wrong:\nprint(f\"x = {x}\")\nprints something like\nTraced<ShapedArray(float32[])>\n, not the number\n5.0\n. During tracing, that's genuinely all\nx\nis.\nif x > 0:\nmay raise a\nConcretizationTypeError\n. JAX can't evaluate a tracer as a Python boolean, because during tracing it doesn't know what the value will actually be at runtime.\npdb.set_trace()\nand\nbreakpoint()\nsuffer the same fate, they can show you a tracer's shape and dtype, but not the concrete numbers flowing through your model at runtime.\nThis isn't a bug in JAX. It's the direct cost of the speed we've been chasing since Week 1. Compilation requires JAX to reason about your function abstractly, once, before any real data exists. The tools we need have to be built for that reality.\njax.debug.print: The JAX-native print()\nThe most direct fix is\njax.debug.print\n. It's designed to be embedded\ninto\nthe compiled computation graph itself, so it prints real, concrete values at actual runtime, even inside\njit\n,\nvmap\n, and\ngrad\n.\nimport jax\nimport jax.numpy as jnp\n\n@jax.jit\ndef debug_forward(x, w):\n    logits = jnp.dot(x, w)\n    jax.debug.print(\"logits = {}\", logits)\n    return jax.nn.softmax(logits)\n\nx = jnp.array([[1.0, 2.0, 3.0]])\nw = jnp.array([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]])\n\noutput = debug_forward(x, w)\nOutput:\nlogits = [[2.2 2.8]]\nThat's a real, concrete array, not a tracer. The print statement got compiled\ninto\nthe XLA graph as an actual print operation, so it fires with real data every time the function runs.\nKeeping output in order\nBecause JAX can reorder and parallelize operations for performance, prints inside heavily parallel code (like\nvmap\n) can come out of sequence. If you need strictly ordered output; for example, watching a value evolve across training steps, use\nordered=True\n:\njax.debug.print(\"step {} loss = {}\", step, loss, ordered=True)\nPrinting inside a training step\nThis is where\njax.debug.print\nearns its keep; dropped directly into a real NNX training step:\nfrom flax import nnx\nimport optax\n\n@nnx.jit\ndef train_step(model, optimizer, batch):\n    def loss_fn(model):\n        logits = model(batch['image'])\n        loss = optax.softmax_cross_entropy_with_integer_labels(\n            logits=logits, labels=batch['label']\n        ).mean()\n        return loss, logits\n    \n    (loss, logits), grads = nnx.value_and_grad(loss_fn, has_aux=True)(model)\n    \n    # Peek at real numbers, right inside the compiled step\n    jax.debug.print(\"loss = {loss}\", loss=loss)\n    jax.debug.print(\"max logit = {m}\", m=jnp.max(logits))\n    \n    optimizer.update(model, grads)\n    return loss\nEvery single call to\ntrain_step\n, compiled or not, will print the real loss and the real max logit. No recompilation needed, no disabling JIT, no performance penalty worth worrying about for occasional prints.\njax.debug.breakpoint: An interactive pause inside JIT\nSometimes a print statement isn't enough. You want to actually stop execution and poke around.\njax.debug.breakpoint()\ngives you that, working like a lightweight version of\npdb\nthat's compatible with tracing.\n@jax.jit\ndef suspicious_function(x, w):\n    hidden = jnp.dot(x, w)\n    \n    jax.debug.breakpoint()  # Execution pauses here at RUNTIME\n    \n    return jax.nn.relu(hidden)\n\noutput = suspicious_function(x, w)\nWhen this runs, you'll drop into an interactive\n(jaxdb)\nprompt:\nEntering jaxdb:\n(jaxdb) p hidden\nArray([[2.2, 2.8]], dtype=float32)\n(jaxdb) c\nThe command set is intentionally limited compared to full\npdb\n:\nCommand\nEffect\np <var>\nPrint a variable's real value\nc\nContinue execution\nq\nQuit\nIt's not a full debugger, but it's enough to answer the question that matters most:\n\"What does this tensor actually look like right now?\"\nMaking it conditional\nYou rarely want to break on every single call; usually you want to break only when something specific goes wrong, like a NaN appearing:\nimport jax.numpy as jnp\nfrom jax import lax\n\n@jax.jit\ndef conditional_breakpoint(x):\n    has_nan = jnp.any(jnp.isnan(x))\n    \n    lax.cond(\n        has_nan,\n        lambda: jax.debug.breakpoint(),\n        lambda: None,\n    )\n    \n    return x * 2\nNow the breakpoint only fires on the exact batch that introduces a NaN — instead of interrupting every single training step.\njax.disable_jit: Bringing back standard Python tools\nSometimes you don't want JAX's limited debugging vocabulary; you want your full, familiar Python debugger back:\npdb\n, your IDE's breakpoints, real\nprint()\nstatements, everything.\njax.disable_jit()\nforces JAX into fully eager execution, exactly like plain NumPy. Every operation runs immediately with real values, so all your normal tools work again.\nimport jax\n\ndef my_function(x, w):\n    hidden = jnp.dot(x, w)\n    print(f\"hidden = {hidden}\")   # Works normally now!\n    breakpoint()                   # Standard Python debugger works!\n    return jax.nn.relu(hidden)\n\nwith jax.disable_jit():\n    output = my_function(x, w)\nYou can also flip this globally, which is convenient while iterating in a notebook:\n# Global flag\njax.config.update(\"jax_disable_jit\", True)\n\n# Or via environment variable, before starting Python\n# JAX_DISABLE_JIT=1 python train.py\nThe trade-off\njax.disable_jit()\ngives you back full debugging power at the cost of everything that makes JAX fast; no fusion, no compilation, no XLA optimization. On a real training step, this can be 10-100x slower.\nTreat it as a scalpel, not your default mode:\n# Debugging: eager execution, full pdb access\nwith jax.disable_jit():\n    debug_output = train_step(model, optimizer, tiny_batch)\n\n# Production: full speed\noutput = train_step(model, optimizer, real_batch)\nOne limitation worth knowing:\ndisable_jit()\ndoesn't always let you step\ninside\nthe internals of\njax.vmap\nor\njax.scan\n: those transformations have their own internal execution logic that isn't simply \"turned off\" by disabling JIT.\nAutomatic NaN detection\nWe spent last week manually hunting NaNs with\nchex.assert_tree_all_finite\n. JAX also has a built-in, lower-level tool for the same problem: the\njax_debug_nans\nflag.\njax.config.update(\"jax_debug_nans\", True)\nWith this enabled, the moment any operation produces a NaN, JAX automatically re-runs that specific operation\neagerly\n, outside of JIT, so it can point you directly at the exact line and operation responsible.\nimport jax\njax.config.update(\"jax_debug_nans\", True)\n\n@jax.jit\ndef buggy_function(x):\n    y = jnp.log(x)  # log(0) = -inf, log(negative) = NaN\n    return y * 2\n\n# This will halt immediately with a traceback pointing at jnp.log\nresult = buggy_function(jnp.array([-1.0, 2.0, 3.0]))\nInstead of your model silently going to NaN somewhere deep inside a 50-layer network and you having no idea which layer did it, JAX interrupts execution the instant the NaN is born and shows you exactly where.\nLimitations\nIt significantly slows down executio; this is a debugging flag, not something to leave on in production.\nIt doesn't work with\njax.shard_map\n.\nIt can occasionally flag NaNs that are actually\nintentional\n(some numerically stable algorithms briefly produce and then discard NaN-like intermediate values), so a small amount of judgment is still required.\nUse it the same way you'd use\nchexify\n: turned on while chasing a bug, turned off before a long production run.\nnnx.display: Seeing your model's structure\nShifting away from runtime values, sometimes the bug isn't in the math; it's in the architecture itself. Maybe a layer has the wrong shape, or a parameter didn't initialize the way you expected.\nnnx.display()\nis the Flax NNX equivalent of\nprint(pytorch_model)\n, but considerably richer:\nfrom flax import nnx\n\nclass CNN(nnx.Module):\n    def __init__(self, *, rngs: nnx.Rngs):\n        self.conv1 = nnx.Conv(1, 32, kernel_size=(3, 3), rngs=rngs)\n        self.conv2 = nnx.Conv(32, 64, kernel_size=(3, 3), rngs=rngs)\n        self.linear = nnx.Linear(3136, 10, rngs=rngs)\n    \n    def __call__(self, x):\n        x = nnx.relu(self.conv1(x))\n        x = nnx.relu(self.conv2(x))\n        x = x.reshape(x.shape[0], -1)\n        return self.linear(x)\n\nmodel = CNN(rngs=nnx.Rngs(0))\nnnx.display(model)\nIn a notebook, this renders as an interactive, expandable tree showing every submodule, every parameter's exact shape and dtype, and how everything nests together. It's the fastest way to confirm your architecture actually matches what you\nthink\nyou built — invaluable when a shape mismatch several layers deep is causing a downstream Chex assertion to fail, and you're not sure which layer is actually at fault.\nModule.sow: Capturing intermediate activations\nSometimes you don't just want to print an intermediate value onc; you want to\ncapture\nit for later inspection, without changing your model's function signature or return value.\nModule.sow()\nlets a layer stash a value as an attribute on itself during the forward pass:\nfrom flax import nnx\n\nclass InspectableCNN(nnx.Module):\n    def __init__(self, *, rngs: nnx.Rngs):\n        self.conv1 = nnx.Conv(1, 32, kernel_size=(3, 3), rngs=rngs)\n        self.linear = nnx.Linear(21632, 10, rngs=rngs)\n    \n    def __call__(self, x):\n        x = self.conv1(x)\n        x = nnx.relu(x)\n        \n        # Stash this activation for later inspection\n        self.sow(nnx.Intermediate, 'conv1_activation', x)\n        \n        x = x.reshape(x.shape[0], -1)\n        return self.linear(x)\n\nmodel = InspectableCNN(rngs=nnx.Rngs(0))\ndummy_input = jnp.ones((4, 28, 28, 1))\noutput = model(dummy_input)\n\n# Access what got \"sown\"\nprint(model.conv1_activation.value.shape)  # (4, 26, 26, 32)\nThis is especially useful for checking whether a specific activation is behaving reasonably, is it saturating at zero (dead ReLUs)? Is its scale exploding layer over layer?, without restructuring your\n__call__\nmethod to return every intermediate value explicitly.\nBy default, calling the module multiple times appends to a tuple rather than overwriting, so you can also use\nsow\nto accumulate values across a sequence (useful inside RNN-style loops).\nPutting it together: A debugging workflow\nHere's a practical order of operations for a real bug, roughly cheapest-to-most-invasive:\n1. Structural sanity check.\nRun\nnnx.display(model)\nfirst. Confirm the architecture is actually what you intended: shapes, layer order, parameter counts.\n2. Static checks.\nMake sure last week's Chex assertions (\nassert_shape\n,\nassert_rank\n,\nassert_type\n) are in place at your function boundaries. Many \"mysterious\" bugs are actually silent broadcasts.\n3. NaN detection.\nIf the loss goes to NaN, turn on\njax_debug_nans\nfirst — let JAX point you directly at the offending operation before you start manually inserting print statements.\n4. Targeted prints.\nOnce you have a rough idea of where the problem lives, add a few\njax.debug.print\ncalls around it to watch real values move through the computation.\n5. Interactive breakpoint.\nIf prints aren't enough, you need to actually poke at a value's shape or compare two tensors interactively, drop in\njax.debug.breakpoint()\n, ideally gated with a\nlax.cond\nso it only fires on the problematic case.\n6. Full eager debugging.\nIf you're truly stuck and need your entire familiar toolchain,\npdb\n, IDE breakpoints, stepping line by line, wrap the smallest possible piece of code in\njax.disable_jit()\nand accept the performance hit temporarily.\n7. Clean up before production.\nStrip out debug prints, disable\njax_debug_nans\n, remove any\ndisable_jit()\ncontext managers. None of these should ship in your real training run.\nExercises\nThe silent tracer\n: Write a JIT-compiled function with a plain\nprint(x)\ninside it. Call it twice with different-shaped inputs. Notice that the print only fires during\ntracing\n(once per unique shape), not on every call; this is a common source of confusion. Then replace it with\njax.debug.print\nand observe that it fires on every call.\nThe NaN hunt\n: Write a function that computes\n1.0 / (x - 5.0)\nfor a batch of values, one of which is exactly\n5.0\n. Turn on\njax_debug_nans\nand run it. Confirm JAX stops you at the exact division.\nConditional breakpoint\n: Take your Week 4 training step and add a\nlax.cond\n-gated\njax.debug.breakpoint()\nthat only triggers when the loss exceeds some threshold (say,\n10.0\n). Feed it a batch designed to produce a large loss and confirm it pauses.\nSow and inspect\n: Add\nself.sow(...)\ncalls after each convolutional block in your Week 3 CNN. After a forward pass, print the mean and standard deviation of each captured activation. Are any of them suspiciously close to zero (dead layer) or huge (exploding activations)?\nQuick reference\nimport jax\nimport jax.numpy as jnp\nfrom jax import lax\n\n# Runtime Printing Inside JIT\njax.debug.print(\"value = {}\", x)\njax.debug.print(\"step {} loss {}\", step, loss, ordered=True)\n\n# Interactive Breakpoint Inside JIT\njax.debug.breakpoint()\n\n# Conditional breakpoint\nlax.cond(has_nan, lambda: jax.debug.breakpoint(), lambda: None)\n\n# Full Eager Execution (standard pdb/print work)\nwith jax.disable_jit():\n    result = my_function(x)\n\njax.config.update(\"jax_disable_jit\", True)  # Global toggle\n\n# Automatic NaN Detection\njax.config.update(\"jax_debug_nans\", True)\n\n# Flax NNX Structural Inspection\nnnx.display(model)\n\n# Capturing Intermediate Values\nself.sow(nnx.Intermediate, 'my_activation', x)\nmodel.my_activation.value  # Access after forward pass\nWhat's next\nWe can now debug both the\nshape\nof our bugs (Chex,\nweek 6\n) and the\nruntime behavior\nof our bugs (this week). But there's a whole different category of problem we haven't touched: what happens when your model is correct, your training loop is correct, and everything is just... slow?\nNext week, we tackle the unglamorous but critical problem of getting data into your model fast enough to keep the GPU fed. We'll introduce\nGrain\n, JAX's purpose-built data loading library, and build a high-performance pipeline that shuffles, augments, and batches data without ever letting your accelerator sit idle.\nNext week\n:\nFeeding the beast: Data pipelines with grain\n#\njit\n#\njax\n#\nerrors",
  "metadata": {
    "title": "When JIT hides your errors: Debugging JAX",
    "description": "Last week, Chex gave us a firewall against shape mismatches and NaN gradients. That firewall catches the bugs that announce themselves; a tensor with the wrong shape, a gradient that explodes into inf",
    "keywords": null,
    "author": null,
    "og:title": "When JIT hides your errors: Debugging JAX",
    "og:description": "Last week, Chex gave us a firewall against shape mismatches and NaN gradients. That firewall catches the bugs that announce themselves; a tensor with the wrong shape, a gradient that explodes into inf",
    "og:url": "https://kambale.dev/debugging-jax",
    "og:site_name": "Wesley Kambale",
    "og:image": "https://cdn.hashnode.com/uploads/covers/61143e31119030192497a888/249d89a6-26e0-47d9-a304-4c73d8fa3c3f.png",
    "og:image:width": "1200",
    "og:image:height": "630",
    "og:type": "article",
    "twitter:card": "summary_large_image",
    "twitter:creator": "https://twitter.com/WesKambale",
    "twitter:title": "When JIT hides your errors: Debugging JAX",
    "twitter:description": "Last week, Chex gave us a firewall against shape mismatches and NaN gradients. That firewall catches the bugs that announce themselves; a tensor with the wrong shape, a gradient that explodes into inf",
    "twitter:image": "https://cdn.hashnode.com/uploads/covers/61143e31119030192497a888/249d89a6-26e0-47d9-a304-4c73d8fa3c3f.png",
    "article:published_time": "2026-07-18T00:12:11.459Z",
    "article:modified_time": "2026-07-29T01:00:32.637Z",
    "article:author": "https://kambale.dev/@kambale"
  },
  "paywall": {
    "listed_paid": false,
    "detected": false,
    "marker": "",
    "public_only": false,
    "policy": "付费/会员内容只保留页面公开可见部分，不绕过付费墙"
  },
  "crawl": {
    "success": true,
    "status_code": 200,
    "error": ""
  },
  "code_blocks": [
    {
      "file": "attachments/487b29700e83/code/body-000.txt",
      "chars": 235,
      "language": "text"
    },
    {
      "file": "attachments/487b29700e83/code/body-001.txt",
      "chars": 290,
      "language": "text"
    },
    {
      "file": "attachments/487b29700e83/code/body-002.txt",
      "chars": 21,
      "language": "text"
    },
    {
      "file": "attachments/487b29700e83/code/body-003.txt",
      "chars": 63,
      "language": "text"
    },
    {
      "file": "attachments/487b29700e83/code/body-004.txt",
      "chars": 614,
      "language": "text"
    },
    {
      "file": "attachments/487b29700e83/code/body-005.txt",
      "chars": 207,
      "language": "text"
    },
    {
      "file": "attachments/487b29700e83/code/body-006.txt",
      "chars": 78,
      "language": "text"
    },
    {
      "file": "attachments/487b29700e83/code/body-007.txt",
      "chars": 247,
      "language": "text"
    },
    {
      "file": "attachments/487b29700e83/code/body-008.txt",
      "chars": 274,
      "language": "text"
    },
    {
      "file": "attachments/487b29700e83/code/body-009.txt",
      "chars": 148,
      "language": "text"
    },
    {
      "file": "attachments/487b29700e83/code/body-010.txt",
      "chars": 206,
      "language": "text"
    },
    {
      "file": "attachments/487b29700e83/code/body-011.txt",
      "chars": 42,
      "language": "text"
    },
    {
      "file": "attachments/487b29700e83/code/body-012.txt",
      "chars": 280,
      "language": "text"
    },
    {
      "file": "attachments/487b29700e83/code/body-013.txt",
      "chars": 501,
      "language": "text"
    },
    {
      "file": "attachments/487b29700e83/code/body-014.txt",
      "chars": 697,
      "language": "text"
    },
    {
      "file": "attachments/487b29700e83/code/body-015.txt",
      "chars": 750,
      "language": "text"
    }
  ],
  "images": [
    {
      "status": "ok",
      "url": "https://cdn.hashnode.com/uploads/covers/61143e31119030192497a888/249d89a6-26e0-47d9-a304-4c73d8fa3c3f.png",
      "bytes": 67064,
      "reused": true,
      "local_path": "attachments/487b29700e83/images/000.png",
      "src": "https://cdn.hashnode.com/uploads/covers/61143e31119030192497a888/249d89a6-26e0-47d9-a304-4c73d8fa3c3f.png",
      "alt": "When JIT hides your errors: Debugging JAX",
      "downloaded": true
    }
  ],
  "audio": [],
  "videos": [],
  "links": [
    {
      "href": "https://kambale.dev/debugging-jax.md",
      "label": "View as Markdown",
      "type": "a"
    },
    {
      "href": "https://hashnode.com/@kambale",
      "label": "W",
      "type": "a"
    },
    {
      "href": "https://twitter.com/WesKambale",
      "label": "",
      "type": "a"
    },
    {
      "href": "https://linkedin.com/in/WesKambale",
      "label": "",
      "type": "a"
    },
    {
      "href": "https://kambale.dev/series/jax-series",
      "label": "Part of series JAX",
      "type": "a"
    },
    {
      "href": "https://kambale.dev/chex",
      "label": "Chex",
      "type": "a"
    },
    {
      "href": "https://kambale.dev/why-jax-the-numpy-you-know-but-faster",
      "label": "week 1",
      "type": "a"
    },
    {
      "href": "https://kambale.dev/tag/jit",
      "label": "# jit",
      "type": "a"
    },
    {
      "href": "https://kambale.dev/tag/jax",
      "label": "# jax",
      "type": "a"
    },
    {
      "href": "https://kambale.dev/tag/errors",
      "label": "# errors",
      "type": "a"
    }
  ],
  "repair_version": 2,
  "repaired_at": "2026-10-14T10:50:46Z",
  "source_url": "https://kambale.dev/debugging-jax",
  "papers": [],
  "actual_video_files": [],
  "body_chars": 17105,
  "expected_body_images": 1,
  "discovery_status": "ok",
  "delivery_status": "ok",
  "missing_reason": "",
  "raw_html": "attachments/487b29700e83/pages/source.html",
  "local_html": "attachments/487b29700e83/pages/article.html",
  "body_base_path": "attachments/487b29700e83/pages/",
  "body_markdown_path": "attachments/487b29700e83/pages/article.md"
}