At commit 0c17ff9, jax-ml/jax has 12,424 analyzed functions, of which 1,404 are in the critical band — and every one of the top five hotspots is ‘fire’ quadrant, meaning structurally complex and actively changing right now. The highest scorer, jaxpr_subcomp in jax/_src/pallas/mosaic/lowering.py, carries a risk score of 22.38, with a cyclomatic complexity of 56 and a fan-out of 35 — and it was modified three days ago. That combination is the definition of live regression risk: an engineer touching this function today is navigating 56 independent execution paths and 35 distinct callees simultaneously. I would start the review here, then work down through _interpret_jaxpr and lower_jaxpr_into_pipelined_module, which share the same file or subsystem and compound the blast radius.
The table below ranks functions by activity-weighted risk — a score that multiplies structural complexity by recent commit frequency. A function that is both hard to understand (high cyclomatic complexity) and actively changing is a higher priority than one that is complex but untouched. CC = cyclomatic complexity (independent execution paths); ND = max nesting depth; FO = fan-out (distinct callees).
Top 5 Hotspots
| Function | File | Risk | CC | ND | FO |
|---|---|---|---|---|---|
jaxpr_subcomp | jax/_src/pallas/mosaic/lowering.py | 22.4 | 56 | 7 | 35 |
_interpret_jaxpr | jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py | 19.7 | 88 | 6 | 44 |
_parse_hlo_new_format | jax/experimental/source_mapper/hlo.py | 19.2 | 48 | 7 | 18 |
_check_jaxpr | jax/_src/core.py | 19.1 | 45 | 5 | 30 |
lower_jaxpr_into_pipelined_module | jax/_src/pallas/mosaic/lowering.py | 19.0 | 72 | 5 | 67 |
Large Repo Analysis
jax is a large repository. To stay within memory constraints, this analysis used hybrid touch mode: structural complexity — CC, ND, FO — is measured precisely for every function. Git activity is tracked at the function level (via git log -L) only for files with 5 or more commits in the last 30 days; other files use a file-level approximation. Rankings therefore surface functions that are both structurally complex and in the most actively-changing parts of the codebase. Dormant code with high structural complexity will rank lower than it would under a full per-function analysis — to surface it, run hotspots analyze . --per-function-touches on a machine with sufficient memory.
JAX is Google’s numerical computation library built around composable function transformations — jit, grad, vmap — and its Pallas extension brings kernel programming to TPU (Mosaic) and GPU (Mosaic GPU) backends. The functions that translate JAX’s internal representation (jaxpr) into hardware-specific code are, structurally speaking, the most demanding in the codebase. That’s exactly where the risk concentrates.
12,424 functions analyzed
Every function in this repository falls into either the ‘fire’ or ‘watch’ quadrant — zero debt, zero dormant functions. That means there’s no low-activity shelter here: even the complex functions are being touched. The 4,272 ‘fire’-quadrant functions are the ones that demand immediate attention.
High cyclomatic complexity — many independent execution paths, each a potential bug surface and required test case.Deeply Nested×5Deeply Nested
Control structures nested 4+ levels deep, making it hard to reason about the full execution state at inner branches.Exit Heavy×5Exit Heavy
Multiple return or throw paths dispersed through the body — each exit needs separate test coverage.God Function×5God Function
Calls an unusually large number of distinct functions (high fan-out), making it the structural centre of gravity for a subsystem.Long Function×5Long Function
Function body is too long to review in a single pass; likely contains multiple distinct responsibilities.Cyclic Hub×1Cyclic Hub
Participates in a call cycle with other high-traffic functions, creating circular dependency risk.Hub Function×1Hub Function
Many other functions call this one — a change here ripples widely through callers.
Every single top-five function is flagged simultaneously for complex branching, deep nesting, multiple exit paths, god-function coupling, and excessive length. That’s not a coincidence — it reflects the inherent difficulty of lowering a general-purpose IR to multiple specialized hardware backends. But difficulty of domain doesn’t reduce the maintenance burden; it amplifies it.
jaxpr_subcomp — jax/_src/pallas/mosaic/lowering.py
By name and path, jaxpr_subcomp handles sub-expression compilation of jaxprs in the Mosaic (TPU) lowering pipeline. The metrics tell a pointed story: cyclomatic complexity of 56 means there are 56 independent execution paths through this function — each one a distinct test case that may or may not exist. A nesting depth of 7 means there are control structures seven levels deep, which makes it genuinely hard to reason about which branch you’re on at any given point without an IDE and significant patience. The fan-out of 35 distinct callees is particularly sharp in Python, where duck typing means those 35 call sites may resolve to different concrete types at runtime — the coupling is broader than the number alone suggests.
This function also carries both the cyclic_hub and hub_function patterns — the only function in the top five to do so. That means it’s not just complex internally; it occupies a central position in the call graph where changes here can create unexpected ripple effects across callers and callees alike. The file has two commits and two authors in the last 90 days, and jaxpr_subcomp was touched three days ago. There’s no historical bug-linked commit signal, which is worth noting: this isn’t a function with a documented defect history, but its structural profile means any change made under time pressure is navigating significant complexity. My recommendation here is to identify the distinct lowering cases embedded in those 56 paths and extract each into a named helper — even pulling out three or four would reduce the effective CC and make the branching structure legible.
_interpret_jaxpr — jax/_src/pallas/mosaic/interpret/interpret_pallas_call.py
_interpret_jaxpr lives in the Pallas Mosaic interpreter — the path that runs a Pallas kernel through a software simulation rather than compiled hardware execution. Interpretation of a general jaxpr requires dispatching on every primitive operation, handling control flow, managing environments, and dealing with all the edge cases that compiled lowering can sometimes elide. A cyclomatic complexity of 88 is the highest in the top five and places this firmly in the “extreme” range — 88 paths through a single function. The nesting depth of 6 and fan-out of 44 compound that: 44 distinct callees in a dynamically typed language means the actual coupling surface is wide and partially invisible until runtime.
The exit_heavy pattern flags multiple return paths, each of which is an implicit contract about what state the caller receives. Combined with 88 branch points, achieving meaningful test coverage here is a significant undertaking. The file has a single commit and a single author in the last 90 days, and no bug-linked commits on record. That single-author concentration is worth tracking: if this interpreter path is modified again, the structural complexity demands that more than one engineer reviews the change. My recommendation is to apply extract-method refactoring primitive-by-primitive — each primitive handler is a natural extraction candidate that would reduce CC by several points per extraction and make the interpreter’s dispatch logic explicit rather than embedded in one monolithic function.
_parse_hlo_new_format — jax/experimental/source_mapper/hlo.py
The experimental/ prefix in the file path and the function’s _new_format suffix are both worth flagging. This is explicitly experimental infrastructure — a source mapper for HLO (High Level Operations, XLA’s IR) — and the _new_format name implies it’s handling a format transition, possibly alongside an older parsing path. A cyclomatic complexity of 48 and nesting depth of 7 in a parser that is not yet stable means the team is actively iterating on something that hasn’t settled. HLO format parsing is inherently branchy — it needs to handle every operation type, every attribute variant, and format-version differences — but CC 48 with ND 7 suggests the format handling has grown organically without structural decomposition.
The god_function and long_function patterns together suggest this is absorbing responsibilities that belong in smaller, focused parsers. No bug-linked commits exist at the file level, and there’s a single author in the last 90 days. Given that this lives under experimental/, I’d prioritize adding a parsing-stage test suite before the next format iteration, then decompose the function by HLO node category. A parser of this complexity warrants explicit grammar-level decomposition — each top-level HLO construct as its own parsing function — rather than a single monolith.
_check_jaxpr — jax/_src/core.py
_check_jaxpr in jax/_src/core.py is the jaxpr validity checker — the function that verifies structural correctness of JAX’s internal program representation before it gets compiled or transformed. This is foundational infrastructure: nearly every JAX transformation path runs through jaxpr validation at some point. A cyclomatic complexity of 45 is consistent with a function that must check every possible jaxpr construct, variable binding rule, and type constraint. The fan-out of 30 reflects the breadth of what it validates — each callee is likely a type predicate, an environment lookup, or a sub-check for a specific jaxpr primitive.
What makes this function higher priority than its raw complexity score might suggest is its position in core.py. This is not an experimental path or a backend-specific lowering — it’s the central validity gate for the entire system. The exit_heavy pattern means there are multiple early-exit paths, which in a checker context likely correspond to different failure modes, each implicitly encoding a type of jaxpr invariant violation. The file has two commits and a single author in the last 90 days, with no bug-linked history. My recommendation here is to decompose the checker by jaxpr node kind — each primitive, each sub-expression type, each binding rule — into individually testable check functions. That both reduces CC and makes invariants explicit and auditable.
lower_jaxpr_into_pipelined_module — jax/_src/pallas/mosaic/lowering.py
This is the second function from jax/_src/pallas/mosaic/lowering.py in the top five, and it’s the one with the most striking single metric in the entire dataset: a fan-out of 67. Sixty-seven distinct functions called from one function. In Python, where argument types are resolved at runtime and decorator stacks can redirect dispatch invisibly, a fan-out of 67 is not 67 dependencies — it’s a lower bound on the actual coupling surface.
By name, this function orchestrates lowering a jaxpr into a pipelined Mosaic module — the TPU’s execution model for overlapping computation and memory transfers. Pipelining inherently involves staging, buffering, and sequencing decisions, which explains some of the complexity, but CC 72 with ND 5 and FO 67 goes well beyond what the domain strictly requires. The god_function pattern here is apt: this function appears to be accumulating pipeline configuration, buffer allocation, and lowering orchestration responsibilities that could be separated. Two authors have touched this file in the last 90 days, and the function was last changed three days ago. Given it shares a file with jaxpr_subcomp, any refactoring in this file should be coordinated — decomposing both functions in isolation risks creating new coupling between the extracted helpers. I’d treat this file as a single refactoring unit and plan the work accordingly.
Patterns Found
Antipatterns detected across the top functions in this snapshot:
| Pattern | Occurrences |
|---|---|
complex_branching | 5 |
deeply_nested | 5 |
exit_heavy | 5 |
god_function | 5 |
long_function | 5 |
cyclic_hub | 1 |
hub_function | 1 |
These labels belong to two tiers — Tier 1 (structural): complex_branching, deeply_nested, exit_heavy, long_function, god_function. Tier 2 (relational/temporal): hub_function, cyclic_hub, middle_man, neighbor_risk, stale_complex, churn_magnet, shotgun_target, volatile_god.
See more analyses with these patterns: complex_branching, deeply_nested, exit_heavy, god_function, long_function.
Reproduce This Analysis
git clone https://github.com/jax-ml/jax
cd jax
git checkout 0c17ff9c866088155e969799766e54925c1ad0dd
hotspots analyze . --mode snapshot --explain-patterns --force --hybrid-touches 5To run the same analysis on your own codebase, run hotspots analyze . --mode snapshot in any local git repo — no configuration required.
I use Hotspots to highlight structural and activity risk — not “bad code.” I treat these findings as a prioritization aid, not a bug predictor. Editorial policy →