From ce3fd0a6788c264b827e36e8fc46a7b4b523218f Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Thu, 3 Sep 2026 20:22:40 -0600
Subject: [PATCH 01/73] A derived-type object allocated many components at
 once, sized by itself

CLUBB's grid: setup_grid allocates every component in one statement,
allocate( gr%zm(ngrdcol,gr%nzm), gr%zt(ngrdcol,gr%nzt), ... ), naming the
object by its dummy rather than this, sizing an axis by another component
of the same object and another by the module's private parameters. The
allocation reader takes any object name and every component of a
statement; a bound spelling obj%comp is bound to that component's flat
name, and the component is carried as an input whether or not the body
reads it; a bound from one is spelled by its upper bound alone; and the
driver's extent is declared once when the subprogram already takes it as a
dummy, as CLUBB's do with ngrdcol.

The read/write scope also learns a sibling's generics
(companion_externals): a call spells zt2zm_api, and a name the scope did
not know as a procedure it counted as a read of data.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/fortran/flatten.py   | 86 ++++++++++++++++++++++++++++-----
 src/recast/fortran/interface.py | 13 +++++
 tests/test_flatten.py           | 83 +++++++++++++++++++++++++++++++
 tests/test_fortran_analysis.py  |  9 ++++
 4 files changed, 178 insertions(+), 13 deletions(-)

diff --git a/src/recast/fortran/flatten.py b/src/recast/fortran/flatten.py
index e6527cc..55ce8b2 100644
--- a/src/recast/fortran/flatten.py
+++ b/src/recast/fortran/flatten.py
@@ -62,7 +62,8 @@
 ]
 
 DERIVED = re.compile(r"UNKNOWN\(TYPE\((\w+)\)\)", re.IGNORECASE)
-ALLOCATE_THIS = re.compile(r"allocate\s*\(\s*this\s*%\s*(\w+)\s*\(([^()]*)\)", re.I)
+ALLOCATE_STMT = re.compile(r"\ballocate\s*\(", re.I)
+COMPONENT_ALLOCATION = re.compile(r"(\w+)\s*%\s*(\w+)\s*\(([^()]*)\)")
 FORTRAN_TYPES = {"float64": "real(8)", "float32": "real(4)", "int32": "integer", "bool": "logical"}
 
 
@@ -218,15 +219,18 @@ def flat_args(self) -> list[dict[str, Any]]:
                         sized.append({"lb": "1", "ub": extent})
                 entry["dims"] = sized
             args.append(entry)
-        args.append(
-            {
-                "name": self.patch_count,
-                "dtype": "int32",
-                "intent": "IN",
-                "optional": False,
-                "dims": None,
-            }
-        )
+        if not any(a["name"].lower() == self.patch_count for a in args):
+            # CLUBB passes ngrdcol explicitly: the driver's extent is a dummy
+            # already there, and declaring it twice does not compile.
+            args.append(
+                {
+                    "name": self.patch_count,
+                    "dtype": "int32",
+                    "intent": "IN",
+                    "optional": False,
+                    "dims": None,
+                }
+            )
         for obj in self.objects:
             for comp in obj.components:
                 args.append(
@@ -317,10 +321,36 @@ def _state_declaration(name: str, root: Path) -> tuple[str, str] | None:
 
 
 def _allocation_bounds(path: Path) -> dict[str, list[str]]:
-    """``component -> [axis bound text, ...]`` from ``allocate (this%c (…))``."""
+    """``component -> [axis bound text, ...]`` from the module's ALLOCATE
+    statements: ``allocate (this%c (…))`` in the CLM family, and CLUBB's
+    ``allocate( gr%zm(ngrdcol,gr%nzm), gr%zt(ngrdcol,gr%nzt), … )`` -- one
+    statement over many components of an object named however the setup
+    routine names its dummy. Continuation lines and comments are folded
+    first; the first allocation of a component wins."""
     out: dict[str, list[str]] = {}
-    for match in ALLOCATE_THIS.finditer(path.read_text(errors="replace")):
-        out.setdefault(match.group(1).lower(), [b.strip() for b in match.group(2).split(",")])
+    text = path.read_text(errors="replace")
+    for start in ALLOCATE_STMT.finditer(text):
+        # The statement: from ``allocate(`` to its matching parenthesis,
+        # across ``&`` continuations, with trailing comments dropped.
+        depth, at = 0, start.end() - 1
+        while at < len(text):
+            character = text[at]
+            if character == "!":
+                at = text.find("\n", at)
+                if at < 0:
+                    break
+                continue
+            if character == "(":
+                depth += 1
+            elif character == ")":
+                depth -= 1
+                if depth == 0:
+                    break
+            at += 1
+        body = text[start.end() : at]
+        for item in COMPONENT_ALLOCATION.finditer(body.replace("&", " ")):
+            axes = [b.strip() for b in item.group(3).split(",")]
+            out.setdefault(item.group(2).lower(), axes)
     return out
 
 
@@ -340,6 +370,11 @@ def spell(text: str) -> str | None:
             return text
         if text in constants:
             return str(constants[text])
+        # A scalar component of the object itself -- CLUBB allocates
+        # ``gr%zm(ngrdcol, gr%nzm)`` -- spelled by that component's flat
+        # name once the plan's objects are known (``_bind_symbolic_extents``).
+        if re.fullmatch(r"\w+\s*%\s*\w+", text):
+            return text
         # An arithmetic bound over constants: ``-nlevsno+1``.
         expression = re.sub(
             r"[A-Za-z_]\w*", lambda m: str(constants.get(m.group(0).lower(), m.group(0))), text
@@ -363,6 +398,8 @@ def spell(text: str) -> str | None:
         extent = patch if low == "1" else None
     elif re.fullmatch(r"-?\d+", low) and re.fullmatch(r"-?\d+", high):
         extent = str(int(high) - int(low) + 1)
+    elif low == "1":
+        extent = high
     else:
         extent = f"({high}) - ({low}) + 1"
     if extent is None:
@@ -802,6 +839,15 @@ def type_info(type_name: str) -> tuple[dict[str, Any], dict[str, list[str]]] | N
                 (type_files[flat.type_name],),
                 kinds,
             )
+            # A component whose allocation is sized by another component of
+            # the same object (``gr%zm(ngrdcol, gr%nzm)``) needs that one
+            # carried too, as an input, whether or not the body reads it.
+            for member, member_axes in bounds.items():
+                if member not in touched[obj]:
+                    continue
+                for ref in re.findall(r"(\w+)\s*%\s*(\w+)", " ".join(member_axes)):
+                    if ref[0].lower() == obj and ref[1].lower() in comps:
+                        touched[obj].setdefault(ref[1].lower(), False)
             for member, written in sorted(touched[obj].items()):
                 spec = comps.get(member)
                 if spec is None:
@@ -972,9 +1018,21 @@ def _bind_symbolic_extents(plan: FlatPlan) -> None:
     answers for makes the component unsupported."""
     by_name = {state.name: state.flat for state in plan.states if not state.extents}
 
+    carried = {(obj.name, comp.name): comp.flat for obj in plan.objects for comp in obj.components}
+    flat_names = set(carried.values())
+
     def bind(text: str) -> str | None:
         missing: list[str] = []
 
+        def component(match: re.Match[str]) -> str:
+            key = (match.group(1).lower(), match.group(2).lower())
+            if key in carried:
+                return carried[key]
+            missing.append(match.group(0))
+            return match.group(0)
+
+        text = re.sub(r"(\w+)\s*%\s*(\w+)", component, text)
+
         def swap(match: re.Match[str]) -> str:
             token = match.group(0)
             lowered = token.lower()
@@ -982,6 +1040,8 @@ def swap(match: re.Match[str]) -> str:
                 return token
             if lowered in by_name:
                 return by_name[lowered]
+            if token in flat_names:
+                return token  # a component of the object, bound just above
             missing.append(token)
             return token
 
diff --git a/src/recast/fortran/interface.py b/src/recast/fortran/interface.py
index cb10bc3..10926bb 100644
--- a/src/recast/fortran/interface.py
+++ b/src/recast/fortran/interface.py
@@ -1305,6 +1305,19 @@ def companion_externals(record: dict[str, Any]) -> dict[str, dict[str, Any]]:
                 if argument["intent"] in ("OUT", "INOUT")
             ],
         }
+    # The sibling's generics too: a call spells the generic (CLUBB's
+    # ``zt2zm_api`` over grid_class's specifics), and a name the scope does
+    # not know as a procedure it counts as a read of data. The entry's
+    # writes are the union over the specifics -- which agree, in every
+    # generic seen so far, on being functions with no OUT argument.
+    for generic, specifics in (record.get("generics") or {}).items():
+        known = [table[s] for s in specifics if s in table]
+        if generic in table or not known:
+            continue
+        table[generic] = {
+            "kind": known[0]["kind"],
+            "out_positions": sorted({at for entry in known for at in entry["out_positions"]}),
+        }
     return table
 
 
diff --git a/tests/test_flatten.py b/tests/test_flatten.py
index 35c75c0..c68ba1f 100644
--- a/tests/test_flatten.py
+++ b/tests/test_flatten.py
@@ -484,3 +484,86 @@ def test_the_adapter_declares_a_lower_bound_and_calls_through_a_generic() -> Non
     assert "use solve_mod, only: solve\n" in text  # both specifics, one generic
     assert "call solve(n=n, x=x)" in text
     assert "solve_one(" not in text.replace("subroutine solve_one_flat(", "")
+
+
+# --- a CLUBB-shaped object: many components in one ALLOCATE, sized by itself
+
+GRID = """\
+module grid_class
+  implicit none
+  private
+  public :: grid, setup_grid
+  integer, parameter :: t_above = 1, t_below = 2
+  type grid
+    integer :: nzm, nzt
+    real(8), allocatable, dimension(:,:) :: zm, zt
+    real(8), allocatable, dimension(:,:,:) :: weights_zt2zm
+    real(8) :: grid_dir
+  end type grid
+contains
+  subroutine setup_grid( ngrdcol, nzmax, gr )
+    integer, intent(in) :: ngrdcol, nzmax
+    type(grid), intent(inout) :: gr
+    integer :: ierr
+    gr%nzm = nzmax
+    gr%nzt = nzmax - 1
+    allocate( gr%zm(ngrdcol,gr%nzm), gr%zt(ngrdcol,gr%nzt), & ! two at once
+              gr%weights_zt2zm(ngrdcol,gr%nzm,t_above:t_below), &
+              stat=ierr )
+    gr%grid_dir = 1.0d0
+  end subroutine setup_grid
+end module grid_class
+"""
+
+COLUMN = """\
+module column_mod
+  use grid_class, only: grid
+  implicit none
+  private
+  public :: ddz
+contains
+  subroutine ddz( nzm, ngrdcol, gr, x, dxdz )
+    integer, intent(in) :: nzm, ngrdcol
+    type(grid), intent(in) :: gr
+    real(8), intent(in), dimension(ngrdcol, nzm) :: x
+    real(8), intent(out), dimension(ngrdcol, nzm) :: dxdz
+    integer :: i, k
+    do k = 1, nzm
+      do i = 1, ngrdcol
+        dxdz(i,k) = gr%grid_dir * x(i,k) * gr%zm(i,k) * gr%weights_zt2zm(i,k,1)
+      end do
+    end do
+  end subroutine ddz
+end module column_mod
+"""
+
+CLUBB_CONVENTIONS = FlatConventions(patch_count="ngrdcol", bounds_pattern=r"^ngrdcol$")
+
+
+def test_an_object_allocated_many_at_once_and_sized_by_itself(tmp_path: Path) -> None:
+    """CLUBB's grid: one ALLOCATE over every component, the object named by
+    the setup routine's dummy rather than ``this``, an axis sized by another
+    component of the same object (``gr%nzm``) and one by the module's
+    private parameters (``t_above:t_below``). The plan carries ``nzm`` as an
+    input the body never reads, spells the extents by it, and declares the
+    driver's extent once, because ``ngrdcol`` is a dummy already."""
+    (tmp_path / "grid_class.f90").write_text(GRID)
+    (tmp_path / "column_mod.f90").write_text(COLUMN)
+    frontend = FortranFrontend(flatten={"patch_count": "ngrdcol", "bounds_pattern": r"^ngrdcol$"})
+    unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:column_mod")
+    facts = frontend.analyze(unit, tmp_path)
+    (plan,) = plans_for(facts, tmp_path, CLUBB_CONVENTIONS)
+    assert plan.usable, plan.unsupported
+    (gr,) = plan.objects
+    by_name = {c.name: c for c in gr.components}
+    assert set(by_name) == {"grid_dir", "zm", "weights_zt2zm", "nzm"}
+    assert by_name["nzm"].written is False
+    assert by_name["zm"].extents == ["ngrdcol", "gr__nzm"]
+    assert by_name["weights_zt2zm"].extents == ["ngrdcol", "gr__nzm", "2"]
+    assert by_name["weights_zt2zm"].bounds[2] == ("1", "2")
+    names = [a["name"] for a in plan.flat_args]
+    assert names.count("ngrdcol") == 1
+    assert names.index("gr__nzm") < names.index("gr__zm")
+    text = fortran_adapter("column_mod", [plan], [])
+    assert "real(8), intent(in) :: gr__zm(ngrdcol, gr__nzm)" in text
+    assert "allocate(gr%zm(1:ngrdcol, 1:gr__nzm))" in text
diff --git a/tests/test_fortran_analysis.py b/tests/test_fortran_analysis.py
index 8fa8e01..4deb514 100644
--- a/tests/test_fortran_analysis.py
+++ b/tests/test_fortran_analysis.py
@@ -1811,3 +1811,12 @@ def test_the_specifics_of_a_public_generic_are_public_through_it(tmp_path: Path)
     assert by_name["solve_many"]["public"] and by_name["solve_many"]["public_via"] == "solve"
     assert not by_name["helper"]["public"] and "public_via" not in by_name["helper"]
     assert record["generics"] == {"solve": ["solve_one", "solve_many"]}
+
+
+def test_a_siblings_generic_is_a_procedure_to_the_read_write_scope(tmp_path: Path) -> None:
+    """A call into a companion spells the generic (CLUBB's ``zt2zm_api``);
+    without an entry for it the scope counted the name as a read of data."""
+    record = interface.extract(_write(tmp_path, "solve.f90", PUBLIC_GENERIC))
+    table = interface.companion_externals(record)
+    assert table["solve"] == {"kind": "subroutine", "out_positions": [1, 2]}
+    assert table["solve_one"]["out_positions"] == [1]

From 7cf2c93baa5c31a63ff572e4c296a3a5cb25f223 Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Thu, 3 Sep 2026 20:31:06 -0600
Subject: [PATCH 02/73] A bound may name any object's component; a sibling's
 buffer OUT is read; a stub under control flow is a stub block

calc_pressure declares thvm(ngrdcol, gr%nzt): a dummy sized by a component
of the object it takes. The flat signature spells that component's flat
name, an argument of the adapter, and the plan carries the component as an
input whether or not the body reads it -- also when another object's
allocation names it (sponge_layer_damping's profile, sized by gr%nzm).
Every dummy's type is read before the objects are walked so the reference
resolves whichever order they come in.

Skx_func(..., Skw_zm) into a sibling: the caller's storage is passed in and
returned, so the caller reads the actual as well as writing it. The
companion table now says which positions are buffers, and the read/write
scope reads them, the rule it already applied to its own subprograms.

if ( stats%l_sample ) then / call stats_update(...) / end if is a
framework stub under a condition. The condition is read on both sides and
the only disagreement is the stub's actuals; the block is waived like a
bare stub, and the verdict names the calls.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/fortran/flatten.py   | 50 ++++++++++++++++++++++++++-------
 src/recast/fortran/interface.py |  9 ++++++
 src/recast/fortran/rwset.py     |  5 +++-
 src/recast/verify/rwset.py      | 10 ++++++-
 tests/test_fortran_analysis.py  |  5 ++--
 5 files changed, 65 insertions(+), 14 deletions(-)

diff --git a/src/recast/fortran/flatten.py b/src/recast/fortran/flatten.py
index 55ce8b2..c69eb1c 100644
--- a/src/recast/fortran/flatten.py
+++ b/src/recast/fortran/flatten.py
@@ -153,6 +153,17 @@ class FlatPlan:
     def name(self) -> str:
         return f"{self.subprogram['name']}_flat"
 
+    def _component_names(self, text: str) -> str:
+        """``obj%comp`` in a bound -> the component's flat name, when the
+        plan carries it; left as written otherwise."""
+        carried = {(o.name, c.name): c.flat for o in self.objects for c in o.components}
+
+        def swap(match: re.Match[str]) -> str:
+            key = (match.group(1).lower(), match.group(2).lower())
+            return carried.get(key, match.group(0))
+
+        return re.sub(r"(\w+)\s*%\s*(\w+)", swap, text)
+
     @property
     def usable(self) -> bool:
         return not self.unsupported and bool(self.objects)
@@ -212,7 +223,11 @@ def flat_args(self) -> list[dict[str, Any]]:
                 for dim in entry["dims"]:
                     if dim.get("ub"):
                         ub = str(dim["ub"]).strip().lower()
-                        sized.append({**dim, "ub": str(self.dim_constants.get(ub, dim["ub"]))})
+                        # ``thvm(ngrdcol, gr%nzt)`` (CLUBB's calc_pressure): a
+                        # dummy sized by a component of the object, spelled
+                        # by that component's flat name, an argument here.
+                        ub = self._component_names(ub)
+                        sized.append({**dim, "ub": str(self.dim_constants.get(ub, ub))})
                     else:
                         counter = f"{self.counter_prefix}{entry['name'].lower()}"
                         extent = counter if counter in names else self.patch_count
@@ -813,6 +828,10 @@ def type_info(type_name: str) -> tuple[dict[str, Any], dict[str, list[str]]] | N
             touched.setdefault(obj, {}).setdefault(member, False)
             if name in writes:
                 touched[obj][member] = True
+        # Every dummy's type first, so a bound in one object's allocation
+        # naming another object's component finds that type on record.
+        for type_name in dummies.values():
+            type_info(type_name)
         for obj in sorted(touched):
             if obj in dummies:
                 flat = FlatObject(name=obj, type_name=dummies[obj], kind="dummy")
@@ -839,15 +858,26 @@ def type_info(type_name: str) -> tuple[dict[str, Any], dict[str, list[str]]] | N
                 (type_files[flat.type_name],),
                 kinds,
             )
-            # A component whose allocation is sized by another component of
-            # the same object (``gr%zm(ngrdcol, gr%nzm)``) needs that one
-            # carried too, as an input, whether or not the body reads it.
-            for member, member_axes in bounds.items():
-                if member not in touched[obj]:
-                    continue
-                for ref in re.findall(r"(\w+)\s*%\s*(\w+)", " ".join(member_axes)):
-                    if ref[0].lower() == obj and ref[1].lower() in comps:
-                        touched[obj].setdefault(ref[1].lower(), False)
+            # A component named by a bound is carried as an input whether
+            # or not the body reads it: by a touched component's allocation
+            # (``gr%zm(ngrdcol, gr%nzm)``), by another object's
+            # (``damping_profile%tau_sponge_damp(gr%nzm)``), or by a dummy's
+            # declaration (``thvm(ngrdcol, gr%nzt)``).
+            named_in_bounds = " ".join(
+                " ".join(type_records[dummies[other]][1].get(member) or [])
+                for other, members in touched.items()
+                if other in dummies and dummies[other] in type_records
+                for member in members
+            )
+            named_in_bounds += " " + " ".join(
+                str(d.get(k) or "")
+                for a in sub["args"]
+                for d in (a.get("dims") or ())
+                for k in ("lb", "ub")
+            )
+            for ref in re.findall(r"(\w+)\s*%\s*(\w+)", named_in_bounds):
+                if ref[0].lower() == obj and ref[1].lower() in comps:
+                    touched[obj].setdefault(ref[1].lower(), False)
             for member, written in sorted(touched[obj].items()):
                 spec = comps.get(member)
                 if spec is None:
diff --git a/src/recast/fortran/interface.py b/src/recast/fortran/interface.py
index 10926bb..ef91ae6 100644
--- a/src/recast/fortran/interface.py
+++ b/src/recast/fortran/interface.py
@@ -1304,6 +1304,12 @@ def companion_externals(record: dict[str, Any]) -> dict[str, dict[str, Any]]:
                 for at, argument in enumerate(sub["args"])
                 if argument["intent"] in ("OUT", "INOUT")
             ],
+            # A buffer OUT is passed in and returned (the caller's storage),
+            # so the caller reads the actual as well as writing it -- the
+            # same rule the scope applies to its own subprograms (#38).
+            "buffer_positions": [
+                at for at, argument in enumerate(sub["args"]) if argument.get("buffer")
+            ],
         }
     # The sibling's generics too: a call spells the generic (CLUBB's
     # ``zt2zm_api`` over grid_class's specifics), and a name the scope does
@@ -1317,6 +1323,9 @@ def companion_externals(record: dict[str, Any]) -> dict[str, dict[str, Any]]:
         table[generic] = {
             "kind": known[0]["kind"],
             "out_positions": sorted({at for entry in known for at in entry["out_positions"]}),
+            "buffer_positions": sorted(
+                {at for entry in known for at in entry.get("buffer_positions", [])}
+            ),
         }
     return table
 
diff --git a/src/recast/fortran/rwset.py b/src/recast/fortran/rwset.py
index ccaa26b..7861b2b 100644
--- a/src/recast/fortran/rwset.py
+++ b/src/recast/fortran/rwset.py
@@ -384,10 +384,13 @@ def call(stmt: Any) -> None:
         if callee is None:
             external = scope.externals.get(name)
             out_positions = set(external.get("out_positions", [])) if external else set()
+            buffers = set(external.get("buffer_positions", [])) if external else set()
             for j, actual in enumerate(actuals):
                 if j in out_positions:
                     write_target(actual)
-                else:
+                if j not in out_positions or j in buffers:
+                    # A buffer OUT of a sibling is read as well as written:
+                    # the emitter passes the caller's storage in (#38).
                     reads.update(expr_reads(actual, scope))
             return
 
diff --git a/src/recast/verify/rwset.py b/src/recast/verify/rwset.py
index 1b29a8e..8867389 100644
--- a/src/recast/verify/rwset.py
+++ b/src/recast/verify/rwset.py
@@ -362,6 +362,11 @@ def walk_stmt(node: ast.stmt) -> None:
 STUB_LINE = re.compile(r"^\s*(?:pass\s*#.*\(infra stub\)|#.*)$")
 
 
+CONTROL_LINE = re.compile(r"^\s*(?:(?:if|elif|for|while)\b[^#]*:|else\s*:|pass)\s*(?:#.*)?$")
+"""A line of control flow with nothing of its own: a condition, a loop
+header, an ``else``, a ``pass``. Around stub markers it is still a stub."""
+
+
 def stubbed_blocks(candidate: Candidate) -> dict[str, str]:
     """``"sub/Bnnn" -> reason`` for every block emitted as stub markers only.
 
@@ -390,7 +395,10 @@ def stubbed_blocks(candidate: Candidate) -> dict[str, str]:
             continue
         body = [ln for ln in lines[span[0] - 1 : span[1]] if ln.strip()]
         stubs = [ln for ln in body if "(infra stub)" in ln]
-        if stubs and all(STUB_LINE.match(ln) for ln in body):
+        if stubs and all(STUB_LINE.match(ln) or CONTROL_LINE.match(ln) for ln in body):
+            # ``if (stats%l_sample) then / call stats_update(...) / end if``
+            # (CLUBB) is a stub under a condition: the condition is read on
+            # both sides, and the only disagreement is the stub's actuals.
             calls = sorted({ln.split("#", 1)[1].split("(")[0].strip() for ln in stubs})
             waived[f"{block['subprogram']}/{block['block']}"] = "framework stub: " + ", ".join(
                 calls
diff --git a/tests/test_fortran_analysis.py b/tests/test_fortran_analysis.py
index 4deb514..fa354ce 100644
--- a/tests/test_fortran_analysis.py
+++ b/tests/test_fortran_analysis.py
@@ -978,7 +978,8 @@ def test_companion_externals_derive_from_the_siblings_record(tmp_path: Path) ->
 """
     record = interface.extract(_write(tmp_path, "sib.f90", sibling), kind_assumptions=KINDS)
     table = interface.companion_externals(record)
-    assert table["qsat_water"] == {"kind": "subroutine", "out_positions": [2, 3]}
+    assert table["qsat_water"]["kind"] == "subroutine"
+    assert table["qsat_water"]["out_positions"] == [2, 3]
 
 
 def _sub_node(tmp_path: Path, name: str):
@@ -1818,5 +1819,5 @@ def test_a_siblings_generic_is_a_procedure_to_the_read_write_scope(tmp_path: Pat
     without an entry for it the scope counted the name as a read of data."""
     record = interface.extract(_write(tmp_path, "solve.f90", PUBLIC_GENERIC))
     table = interface.companion_externals(record)
-    assert table["solve"] == {"kind": "subroutine", "out_positions": [1, 2]}
+    assert table["solve"] == {"kind": "subroutine", "out_positions": [1, 2], "buffer_positions": []}
     assert table["solve_one"]["out_positions"] == [1]

From 19b5b2d98bbf1457fb5df2bb2ea527b3362c30fc Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Thu, 3 Sep 2026 20:39:24 -0600
Subject: [PATCH 03/73] Following generics into every specific, dummy-sized
 allocations, private state left to its module, ungated means not sampled

An object passed to a sibling through a generic (zt2zm_api over
grid_class's specifics) was not followed, so the adapter built a record
missing the components the specific reads. A generic is followed into
every specific; the union is what the adapter carries.

An allocation sized by the allocating routine's dummy
(coef_wp4_implicit(1:ngrdcol,1:nz)) is spelled by the planned
subprogram's dummy of the same name, an argument of the adapter.

A module variable the module keeps private (error_code's
clubb_debug_level) reaches no use statement, so no adapter sets it; both
sides run with the module's default and the plan records it under
left_to_module instead of failing.

A subprogram the operator declared ungated is not sampled by the
bit-exact gate either: the declaration says the reference cannot be held
on generated inputs, and rcm_sat_adj proves it by error-stopping the run.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/fortran/flatten.py | 43 +++++++++++++++++++++++++++++++----
 src/recast/verify/bitexact.py | 11 ++++++++-
 2 files changed, 49 insertions(+), 5 deletions(-)

diff --git a/src/recast/fortran/flatten.py b/src/recast/fortran/flatten.py
index c69eb1c..06682d0 100644
--- a/src/recast/fortran/flatten.py
+++ b/src/recast/fortran/flatten.py
@@ -146,6 +146,11 @@ class FlatPlan:
     dim_constants: dict[str, int] = field(default_factory=dict)
     """Named extents of the original dummies that are tree constants
     (``a(nrk,nrk)``), so the flat signature can spell them as numbers."""
+    left_to_module: list[str] = field(default_factory=list)
+    """Module state the body reaches that the adapter cannot set: a
+    variable its module keeps private (CLUBB's ``error_code %
+    clubb_debug_level``). Both sides run with the module's own default,
+    and the plan says so rather than failing."""
     patch_count: str = "np_"
     counter_prefix: str = "num_"
 
@@ -195,6 +200,7 @@ def from_dict(cls, data: dict[str, Any]) -> FlatPlan:
             unsupported=list(data.get("unsupported", [])),
             states=states,
             dim_constants=dict(data.get("dim_constants", {})),
+            left_to_module=list(data.get("left_to_module", [])),
             patch_count=data.get("patch_count", "np_"),
             counter_prefix=data.get("counter_prefix", "num_"),
         )
@@ -583,13 +589,25 @@ def take(stmt: Any) -> None:
         interesting = objects | set(aliases)
         # ``dummy = hybrid(...)`` parses as any of three node kinds depending
         # on what fparser could tell about the name.
+        # A call spelled by a generic (CLUBB's ``zt2zm_api`` over grid_class's
+        # specifics) is followed into every specific: which one Fortran
+        # picks depends on ranks this walk does not resolve, and the union
+        # of what they touch is what the adapter has to carry.
+        generics: dict[str, list[str]] = {}
+        for _path, module_record_, _sub in procedures.values():
+            for generic, specifics in (module_record_.get("generics") or {}).items():
+                generics.setdefault(generic.lower(), [n.lower() for n in specifics])
+        expanded: list[tuple[Any, str]] = []
         for call in [
             *walk(node, f03.Call_Stmt),
             *walk(node, f03.Part_Ref),
             *walk(node, f03.Function_Reference),
             *walk(node, f03.Structure_Constructor),
         ]:
-            callee = str(call.children[0]).lower()
+            spelled_callee = str(call.children[0]).lower()
+            for callee in generics.get(spelled_callee, [spelled_callee]):
+                expanded.append((call, callee))
+        for call, callee in expanded:
             if callee not in procedures:
                 continue
             path, module_record, callee_record = procedures[callee]
@@ -991,12 +1009,22 @@ def _state_vars(
         if not record:
             continue
         module = str(record.get("module", "")).lower()
+        public = {str(n).lower() for n in record.get("public") or ()}
+        default_private = bool(
+            re.search(r"^\s*private\s*(?:!.*)?$", path.read_text(errors="replace"), re.I | re.M)
+        )
         for entry in record.get("module_state", ()):
             name = str(entry["name"]).lower()
             # Every module *variable* the run may have set -- initialized in
             # the tree or not -- is the run's to say; parameters are not here.
             if name not in wanted or name in seen:
                 continue
+            if default_private and name not in public:
+                # No ``use`` reaches it, so no adapter sets it: both sides
+                # run with the module's own default, and the plan says so.
+                seen.add(name)
+                plan.left_to_module.append(f"{module}%{name}")
+                continue
             if (
                 DERIVED.match(str(entry.get("dtype")))
                 or str(entry.get("dtype")) not in FORTRAN_TYPES
@@ -1047,7 +1075,14 @@ def _bind_symbolic_extents(plan: FlatPlan) -> None:
     state name, which is a scalar argument of the adapter; a name no state
     answers for makes the component unsupported."""
     by_name = {state.name: state.flat for state in plan.states if not state.extents}
-
+    # A dummy of the subprogram is an argument of the adapter, so an
+    # allocation sized by the allocating routine's dummy of the same name
+    # (``coef_wp4_implicit(1:ngrdcol,1:nz)``, CLUBB) is spelled as written.
+    dummies = {
+        str(a["name"]).lower()
+        for a in plan.subprogram["args"]
+        if not a.get("dims") and not DERIVED.match(str(a["dtype"]))
+    }
     carried = {(obj.name, comp.name): comp.flat for obj in plan.objects for comp in obj.components}
     flat_names = set(carried.values())
 
@@ -1070,8 +1105,8 @@ def swap(match: re.Match[str]) -> str:
                 return token
             if lowered in by_name:
                 return by_name[lowered]
-            if token in flat_names:
-                return token  # a component of the object, bound just above
+            if token in flat_names or lowered in dummies:
+                return token  # a component bound just above, or a dummy
             missing.append(token)
             return token
 
diff --git a/src/recast/verify/bitexact.py b/src/recast/verify/bitexact.py
index 8992ebe..b13c33b 100644
--- a/src/recast/verify/bitexact.py
+++ b/src/recast/verify/bitexact.py
@@ -276,8 +276,17 @@ def generable(name: str) -> bool:
             skipped = sorted(set(offered) - set(wanted))
         else:
             by_subprogram = {}
+            # One the operator declared ungated is not sampled either: the
+            # declaration says the reference cannot be held on generated
+            # inputs (CLUBB's rcm_sat_adj iterates and error-stops on them).
+            declared_ungated = set(config.get("ungated") or {})
             wanted = config.get("subprograms") or [
-                name for name in wrappers if name in table and judged(name) and generable(name)
+                name
+                for name in wrappers
+                if name in table
+                and judged(name)
+                and generable(name)
+                and name not in declared_ungated
             ]
             skipped = sorted(set(wrappers) - set(wanted))
 

From a5b2391546e1eddd1526b81c9de972e6491a060e Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Thu, 3 Sep 2026 20:50:28 -0600
Subject: [PATCH 04/73] A generic is followed into the specifics its arity
 fits, and the recursion guard is a stack

Following every specific with the caller's positional actuals mapped the
object onto whichever dummy sat at that position -- zm_min for the
one-level specific -- and the specific the call actually reaches was then
skipped, because a nested follow had marked it visited first. The
specifics whose argument count fits the call are followed (the exact
count when any matches), and a key is released when its follow returns:
the set guards against cycles, it is not a memo.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/fortran/flatten.py | 22 ++++++++++++++++++++--
 1 file changed, 20 insertions(+), 2 deletions(-)

diff --git a/src/recast/fortran/flatten.py b/src/recast/fortran/flatten.py
index 06682d0..d3dde44 100644
--- a/src/recast/fortran/flatten.py
+++ b/src/recast/fortran/flatten.py
@@ -605,8 +605,22 @@ def take(stmt: Any) -> None:
             *walk(node, f03.Structure_Constructor),
         ]:
             spelled_callee = str(call.children[0]).lower()
-            for callee in generics.get(spelled_callee, [spelled_callee]):
-                expanded.append((call, callee))
+            candidates = generics.get(spelled_callee)
+            if candidates is None:
+                expanded.append((call, spelled_callee))
+                continue
+            # The specifics a call of this arity can reach: Fortran picks by
+            # rank too, which this walk does not resolve, so the union of
+            # what the arity-compatible ones touch is what the adapter carries.
+            given = len(call.children[1].children) if call.children[1] is not None else 0
+            arity = {
+                specific: len(procedures[specific][2]["args"])
+                for specific in candidates
+                if specific in procedures
+            }
+            exact = [n for n, count in arity.items() if count == given]
+            fitting = exact or [n for n, count in arity.items() if count >= given]
+            expanded.extend((call, callee) for callee in fitting)
         for call, callee in expanded:
             if callee not in procedures:
                 continue
@@ -682,6 +696,7 @@ def take(stmt: Any) -> None:
             visited.add(key)
             callee_node = _subprogram_node(path, callee)
             if callee_node is None:
+                visited.discard(key)
                 continue
             _, inner_reads, inner_writes = _accesses(
                 callee_node,
@@ -695,6 +710,9 @@ def take(stmt: Any) -> None:
                 externals,
                 companions,
             )
+            # A guard against cycles, not a memo: the same callee followed
+            # from another site, under another mapping, is followed again.
+            visited.discard(key)
             callee_dummies = {a["name"].lower() for a in callee_record["args"]}
             for inner, target in ((inner_reads, reads), (inner_writes, writes)):
                 for item in inner:

From 76281d8fb8c35c3ff87a6983f8d2843f8d3e8fe3 Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Thu, 3 Sep 2026 21:08:12 -0600
Subject: [PATCH 05/73] A quotient of real parameters is a real quotient; an
 intent(out) object is read off what the translation returned; a translated
 error stop is a failed comparison

CLUBB's constants_clubb declares ep = Rd / Rv over two real parameters.
The fold's rule -- an initializer with no real literal is integer
arithmetic throughout -- made it zero, and every unit reading ep1 or ep2
computed NaN. The resolver now records each constant's declared base
type, the fold types a name by it, and both renderers carry the
environment along in dependency order; nrk = runge_kutta_type / 10 stays 4.

The Python flat adapter read a written component off the record it built
from the inputs even when the dummy was intent(out) and the translation
returned a new object (sponge_layer_damping's profile): it now rebinds to
the returned object first.

A translated bare error stop raises SystemExit, which the bit-exact gate
let through and which ended the run with no report; it is now a
comparison that could not be made, named as such.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/fortran/expr.py              | 30 ++++++++++++++-----------
 src/recast/fortran/tree.py              |  6 +++--
 src/recast/fortran/use.py               | 14 ++++++++----
 src/recast/transform/numpy/constants.py | 13 ++++++-----
 src/recast/transform/numpy/flat.py      |  6 +++++
 src/recast/verify/bitexact.py           |  6 +++++
 tests/test_fortran_analysis.py          | 29 ++++++++++++++++++++++++
 7 files changed, 80 insertions(+), 24 deletions(-)

diff --git a/src/recast/fortran/expr.py b/src/recast/fortran/expr.py
index 6403575..f3f3aaf 100644
--- a/src/recast/fortran/expr.py
+++ b/src/recast/fortran/expr.py
@@ -205,27 +205,28 @@ def render(
 REAL_CALLS = frozenset({"real", "dble", "sqrt"} | KIND_INQUIRIES)
 
 
-def typed(expr: Expr) -> str | None:
+def typed(expr: Expr, env: dict[str, str | None] | None = None) -> str | None:
     """``"real"``, ``"int"``, or ``None`` when a bare name leaves it open.
 
     Type inference the fold needs for exactly one decision: whether a ``/``
     is Fortran's integer division. A real literal or a real-valued call
     anywhere in an operand makes the quotient real; ``int(...)`` and integer
-    literals make it integer; a name is whatever its initializer was, which
-    this tree does not carry.
+    literals make it integer; a name is what ``env`` says its declaration
+    was -- CLUBB's ``ep = Rd / Rv`` over two real parameters is a real
+    quotient, and a fold that guessed integer made it zero.
     """
     if expr.kind == "real":
         return "real"
     if expr.kind == "int":
         return "int"
     if expr.kind == "name":
-        return None
+        return (env or {}).get(expr.text)
     if expr.kind == "call":
         if expr.text == "int":
             return "int"
         if expr.text in REAL_CALLS:
             return "real"
-    kinds = {typed(a) for a in expr.args}
+    kinds = {typed(a, env) for a in expr.args}
     if "real" in kinds:
         return "real"
     if kinds == {"int"}:
@@ -233,24 +234,27 @@ def typed(expr: Expr) -> str | None:
     return None
 
 
-def with_integer_division(expr: Expr, *, default_integer: bool | None = None) -> Expr:
+def with_integer_division(
+    expr: Expr, *, default_integer: bool | None = None, env: dict[str, str | None] | None = None
+) -> Expr:
     """The tree with every integer ``/`` spelled ``//``.
 
     Fortran divides two integers to an integer: ``nrk = runge_kutta_type / 10``
     is 4, not 4.1. A quotient whose operands are both known integers is
-    marked; one with a name in it falls back to ``default_integer``, which
-    the caller sets from the whole initializer -- an expression with no real
-    literal and no real-valued call in it is integer arithmetic throughout,
-    because a name in it is an integer parameter or it would have had one.
+    marked; one with a name in it is typed by ``env`` (the declared types
+    of the constants resolved so far) and otherwise falls back to
+    ``default_integer``, which the caller sets from the whole initializer.
     """
     if default_integer is None:
-        default_integer = typed(expr) != "real"
+        default_integer = typed(expr, env) != "real"
     if not expr.args:
         return expr
-    args = tuple(with_integer_division(a, default_integer=default_integer) for a in expr.args)
+    args = tuple(
+        with_integer_division(a, default_integer=default_integer, env=env) for a in expr.args
+    )
     text = expr.text
     if expr.kind == "binary" and expr.text == "/":
-        kinds = {typed(a) for a in args}
+        kinds = {typed(a, env) for a in args}
         if kinds == {"int"} or ("real" not in kinds and default_integer):
             text = "//"
     return Expr(expr.kind, text, args)
diff --git a/src/recast/fortran/tree.py b/src/recast/fortran/tree.py
index 1068cc9..c104922 100644
--- a/src/recast/fortran/tree.py
+++ b/src/recast/fortran/tree.py
@@ -159,18 +159,19 @@ def _evaluate(
     initialize it with something a parameter can be folded from."""
     # Lazy, like ``render`` above: ``expr`` parses, and this module is imported
     # by paths that must stay importable without the ``fortran`` extra.
-    from recast.fortran.expr import python_call, with_integer_division
+    from recast.fortran.expr import python_call, typed, with_integer_division
 
     try:
         records = resolve([name], files)
     except unresolved:
         return None
     env: dict[str, Any] = {}
+    kinds: dict[str, str | None] = {}
     try:
         for entry in records:
             # Integer arithmetic where Fortran's ``/`` truncates.
             text = render(
-                with_integer_division(entry["expr"]),
+                with_integer_division(entry["expr"], env=kinds),
                 real=lambda t: f"float('{t}')",
                 integer=lambda t: t,
                 name=lambda t: t.upper(),
@@ -179,6 +180,7 @@ def _evaluate(
             scope = {"__builtins__": {}, "max": max, "min": min, "abs": abs, "int": int}
             scope.update({"float": float, "math": math, "sys": sys})
             env[entry["name"].upper()] = eval(text, scope, dict(env))  # noqa: S307
+            kinds[entry["name"]] = entry.get("dtype") or typed(entry["expr"], kinds)
     except Exception:  # an initializer shape the renderer has no rule for
         return None
     return env.get(name.upper())
diff --git a/src/recast/fortran/use.py b/src/recast/fortran/use.py
index 42f64c5..e2424de 100644
--- a/src/recast/fortran/use.py
+++ b/src/recast/fortran/use.py
@@ -34,7 +34,7 @@ def harvest(path: Path) -> dict[str, tuple[Any, int | None]]:
     of them are harvested.
     """
     ast = parse(path)
-    out: dict[str, tuple[Any, int | None]] = {}
+    out: dict[str, tuple[Any, int | None, str | None]] = {}
     for mod in walk(ast, f03.Module):
         spec = next((c for c in mod.children if isinstance(c, f03.Specification_Part)), None)
         if spec is None:
@@ -46,9 +46,14 @@ def harvest(path: Path) -> dict[str, tuple[Any, int | None]]:
                 if item is not None and getattr(item, "span", None):
                     line = item.span[0]
                     break
+            # The declared base type, which is what says whether ``rd / rv``
+            # is a real quotient: the fold cannot tell from two names.
+            base = str(decl.children[0]).split("(")[0].strip().upper()
+            declared = {"REAL": "real", "DOUBLE PRECISION": "real", "INTEGER": "int"}.get(base)
             for ent in walk(decl, f03.Entity_Decl):
                 if ent.children[3] is not None:
-                    out[str(ent.children[0]).lower()] = (ent.children[3].children[1], line)
+                    initializer = ent.children[3].children[1]
+                    out[str(ent.children[0]).lower()] = (initializer, line, declared)
     return out
 
 
@@ -63,7 +68,7 @@ def resolve(symbols: list[str], sources: list[Path]) -> list[dict[str, Any]]:
     constant that silently becomes undefined downstream is far more expensive to
     diagnose than a failure here that names it.
     """
-    table: dict[str, tuple[Any, int | None]] = {}
+    table: dict[str, tuple[Any, int | None, str | None]] = {}
     origin: dict[str, Path] = {}
     for path in sources:
         for name, rec in harvest(path).items():
@@ -80,7 +85,7 @@ def need(name: str) -> None:
         if name not in table:
             raise UnresolvedConstant(f"no initializer for {name!r} in {[str(s) for s in sources]}")
         seen.add(name)
-        node, line = table[name]
+        node, line, declared = table[name]
         # The one legal self-reference, a kind inquiry on the constant being
         # declared, stands for its kind alone; see ``expr.substitute``.
         expr: Expr = substitute(build(node), name, Expr("real", "1.0"))
@@ -90,6 +95,7 @@ def need(name: str) -> None:
             {
                 "name": name,
                 "expr": expr,
+                "dtype": declared,
                 "source": str(origin[name]),
                 "line": line,
                 "requested": name in requested,
diff --git a/src/recast/transform/numpy/constants.py b/src/recast/transform/numpy/constants.py
index 72f36d9..44f4290 100644
--- a/src/recast/transform/numpy/constants.py
+++ b/src/recast/transform/numpy/constants.py
@@ -23,7 +23,7 @@
 from pathlib import PurePath, PurePosixPath
 from typing import Any
 
-from recast.fortran.expr import Expr, python_call, render, with_integer_division
+from recast.fortran.expr import Expr, python_call, render, typed, with_integer_division
 
 __all__ = [
     "constants_module",
@@ -285,18 +285,21 @@ def use_constants_module(resolved: list[dict[str, Any]], module_name: str) -> st
         "import numpy as np",
         "",
     ]
+    env: dict[str, str | None] = {}
     for entry in resolved:
-        value = _python(entry["expr"])
+        value = _python(entry["expr"], env)
+        env[entry["name"]] = entry.get("dtype") or typed(entry["expr"], env)
         where = f"{PurePath(entry['source']).name}:{entry['line']}"
         lines.append(f"{entry['name'].upper()} = {value}  # {where}")
     return "\n".join(lines) + "\n"
 
 
-def _python(expr: Expr) -> str:
+def _python(expr: Expr, env: dict[str, str | None] | None = None) -> str:
     # Fortran divides two integers to an integer; ``with_integer_division``
-    # spells those quotients ``//`` from the tree's own types.
+    # spells those quotients ``//`` from the tree's own types and the
+    # declared types of the constants before it.
     return render(
-        with_integer_division(expr),
+        with_integer_division(expr, env=env),
         real=lambda text: f"np.float64('{text}')",
         integer=lambda text: text,
         name=lambda text: text.upper(),
diff --git a/src/recast/transform/numpy/flat.py b/src/recast/transform/numpy/flat.py
index 2404762..b82add5 100644
--- a/src/recast/transform/numpy/flat.py
+++ b/src/recast/transform/numpy/flat.py
@@ -77,6 +77,12 @@ def python_adapter(plans: list[FlatPlan]) -> str:
             lines.append(f"    {', '.join(_py(n) + '_' for n in outs)}, = _out")
         for obj in plan.objects:
             if obj.kind == "dummy":
+                if obj.name in outs and obj.name not in {a["name"] for a in passed}:
+                    # An intent(out) object is one the translation returns
+                    # (sponge_layer_damping's profile, CLUBB): its written
+                    # components are read off the returned object, not the
+                    # record built from the inputs.
+                    lines.append(f"    {obj.name} = {_py(obj.name)}_")
                 for comp in obj.components:
                     if comp.written:
                         lines.append(f"    {comp.flat} = {obj.name}.{comp.name}")
diff --git a/src/recast/verify/bitexact.py b/src/recast/verify/bitexact.py
index b13c33b..d00f118 100644
--- a/src/recast/verify/bitexact.py
+++ b/src/recast/verify/bitexact.py
@@ -709,6 +709,12 @@ def _compare_subprogram(
             truth_args = [truth_kwargs[spell(a["name"])] for a in required if a["intent"] != "OUT"]
             try:
                 translated_out = translated_fn(**translated_kwargs)
+            except SystemExit as error:
+                # A translated ``error stop`` (CLUBB's calc_mixture_fraction on
+                # a zero F_x): the candidate's answer for these inputs is to
+                # end the program, which is a comparison that cannot be made,
+                # not the end of the run.
+                return {"error": f"candidate raised: SystemExit (error stop): {error}"}
             except Exception as error:
                 return {"error": f"candidate raised: {type(error).__name__}: {error}"}
             if samples is not None:
diff --git a/tests/test_fortran_analysis.py b/tests/test_fortran_analysis.py
index fa354ce..b0ab4f6 100644
--- a/tests/test_fortran_analysis.py
+++ b/tests/test_fortran_analysis.py
@@ -1821,3 +1821,32 @@ def test_a_siblings_generic_is_a_procedure_to_the_read_write_scope(tmp_path: Pat
     table = interface.companion_externals(record)
     assert table["solve"] == {"kind": "subroutine", "out_positions": [1, 2], "buffer_positions": []}
     assert table["solve_one"]["out_positions"] == [1]
+
+
+def test_a_quotient_of_real_parameters_is_a_real_quotient(tmp_path: Path) -> None:
+    """CLUBB's ``ep = Rd / Rv``: no literal in sight, two real parameters.
+    The old rule -- no real literal means integer arithmetic throughout --
+    folded it to zero. The declared types decide."""
+    import numpy as np
+
+    from recast.transform.numpy.constants import use_constants_module
+
+    src = """\
+module gas_mod
+  implicit none
+  integer, parameter :: core_rknd = 8
+  real( kind = core_rknd ), parameter :: rd = 287.04_core_rknd, rv = 461.5_core_rknd
+  real( kind = core_rknd ), parameter :: ep = rd / rv
+  real( kind = core_rknd ), parameter :: ep2 = 1.0_core_rknd / ep
+  integer, parameter :: runge_kutta_type = 45
+  integer, parameter :: nrk = runge_kutta_type / 10
+end module gas_mod
+"""
+    _write(tmp_path, "gas.f90", src)
+    resolved = use.resolve(["ep2", "nrk"], [tmp_path / "gas.f90"])
+    assert {r["name"]: r["dtype"] for r in resolved}["ep"] == "real"
+    scope: dict[str, object] = {}
+    exec(use_constants_module(resolved, "gas_mod"), scope)  # generated text under test
+    assert scope["EP"] == np.float64(287.04) / np.float64(461.5)
+    assert scope["EP2"] == np.float64(1.0) / scope["EP"]
+    assert scope["NRK"] == 4

From bff9cb45a489b4692763187560966c354e34c14c Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Thu, 3 Sep 2026 21:28:27 -0600
Subject: [PATCH 06/73] A loop index read after its loop has Fortran's
 completion value

CLUBB's lscale_width_vert_avg searches with do k_avg_upper = k, ...; if
(...) exit; end do and integrates up to k_avg_upper. On completion Fortran
leaves the index one step past the end, m1 + n * m3; Python's for leaves
the last value, and the integral ran one level short -- a NaN on one side
in wp23_term_splat_lhs. For a loop whose index some later statement of the
subprogram names, the emitter adds a for-else that sets the completion
value; an EXIT is a break and skips it, as Fortran keeps the exit value.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/numpy/statements.py | 48 ++++++++++++++++++-
 tests/test_numpy_translate.py            | 60 ++++++++++++++++++++++++
 2 files changed, 107 insertions(+), 1 deletion(-)

diff --git a/src/recast/transform/numpy/statements.py b/src/recast/transform/numpy/statements.py
index bb01ead..fbd22b6 100644
--- a/src/recast/transform/numpy/statements.py
+++ b/src/recast/transform/numpy/statements.py
@@ -131,6 +131,33 @@ def derived_array(type_name: str, extents: list[str], known: dict[str, Any]) ->
 write using anything else is refused rather than silently list-directed."""
 
 
+def _loops_whose_index_is_read_after(subprogram: Any) -> set[int]:
+    """The DO constructs whose loop variable some later statement of the
+    subprogram names -- read after the loop, where Fortran's completion
+    value (one step past the end) and Python's (the last value) differ."""
+    from recast.fortran.interface import names_in, node_span
+
+    marked: set[int] = set()
+    for loop in walk(subprogram, (f03.Block_Nonlabel_Do_Construct, f03.Block_Label_Do_Construct)):
+        do_statement = walk(loop, (f03.Nonlabel_Do_Stmt, f03.Label_Do_Stmt))
+        control = walk(do_statement[0], f03.Loop_Control) if do_statement else []
+        if not control or control[0].children[1] is None:
+            continue
+        variable = str(control[0].children[1][0]).lower()
+        _, end_line = node_span(loop)
+        if end_line is None:
+            continue
+        for statement in walk(subprogram):
+            item = getattr(statement, "item", None)
+            span = getattr(item, "span", None) if item is not None else None
+            if not span or span[0] <= end_line:
+                continue
+            if variable in names_in(statement):
+                marked.add(id(loop))
+                break
+    return marked
+
+
 @dataclass
 class Statements:
     """Render Fortran statements for one subprogram.
@@ -177,6 +204,13 @@ class Statements:
     called_names: set[str] = field(default_factory=set)
     """Names this subprogram's body calls or subscripts. Filled by ``scan``."""
 
+    index_read_after: set[int] = field(default_factory=set)
+    """DO constructs (by node id) whose index a later statement reads.
+    Filled by ``scan``. Fortran leaves a completed loop's index one step past
+    the end; Python's ``for`` leaves the last value, so these loops get an
+    ``else`` that sets the completion value (an EXIT, a ``break``, skips it,
+    as Fortran keeps the exit value)."""
+
     exit_labels: dict[int, str] = field(default_factory=dict)
     """``id(do-construct)`` -> the label that means ``exit`` inside it."""
 
@@ -217,6 +251,7 @@ def scan(self, subprogram: Any) -> None:
         """
         self.exit_labels = {}
         self.consumed_labels = set()
+        self.index_read_after = _loops_whose_index_is_read_after(subprogram)
         self.assigned_names = {
             str(a.children[0]).lower()
             for a in walk(subprogram, f03.Assignment_Stmt)
@@ -1392,7 +1427,18 @@ def _do_construct_inner(
                     f"{pad}for {name} in range({low}, "
                     f"({high}) + (1 if ({step}) > 0 else -1), {step}):"
                 )
-        return [head, *self._loop_body(node, indent, cycle_name)]
+        lines = [head, *self._loop_body(node, indent, cycle_name)]
+        if id(node) in self.index_read_after:
+            # CLUBB's lscale_width_vert_avg searches with ``do k_avg_upper =
+            # k, ...; if (...) exit; end do`` and then integrates up to
+            # k_avg_upper: on completion Fortran's index is the first value
+            # past the end, m1 + n * m3, and ``for``'s is the last one.
+            # ``else`` runs exactly when no ``break`` did.
+            increment = step if step is not None else "1"
+            trips = f"max(0, (({high}) - ({low}) + ({increment})) // ({increment}))"
+            lines.append(f"{pad}else:")
+            lines.append(f"{pad}    {name} = ({low}) + {trips} * ({increment})")
+        return lines
 
     def _caught_cycle(self, body: list[str], indent: int, cycle_name: str | None) -> list[str]:
         """A loop body, wrapped so a CYCLE naming *this* loop reaches its header."""
diff --git a/tests/test_numpy_translate.py b/tests/test_numpy_translate.py
index c31bf9b..6151e2a 100644
--- a/tests/test_numpy_translate.py
+++ b/tests/test_numpy_translate.py
@@ -289,3 +289,63 @@ def test_logical_operators_on_arrays_are_elementwise(tmp_path: Path) -> None:
     assert " & " in module
     assert "not l_scalar" in module
     assert "not l_ok" not in module
+
+
+SEARCH_LOOP = """\
+module search_mod
+  implicit none
+  private
+  public :: first_above
+contains
+  subroutine first_above( n, z, zmax, k_found, k_scan )
+    integer, intent(in) :: n
+    real, dimension(n), intent(in) :: z
+    real, intent(in) :: zmax
+    integer, intent(out) :: k_found, k_scan
+    integer :: k, kk
+    do k = 1, n
+      if ( z(k) > zmax ) exit
+    end do
+    k_found = k
+    do kk = 1, n, 2
+      k_scan = kk
+    end do
+    k_scan = kk
+  end subroutine first_above
+end module search_mod
+"""
+
+
+def test_a_loop_index_read_after_the_loop_has_the_completion_value(tmp_path: Path) -> None:
+    """CLUBB's lscale_width_vert_avg searches with ``do k = ...; if (...)
+    exit; end do`` and integrates up to ``k`` afterwards. On completion
+    Fortran leaves the index one step past the end -- ``n + 1`` for a unit
+    step, the first odd value past ``n`` for a step of two -- and after an
+    EXIT it keeps the exit value. Python's ``for`` leaves the last value."""
+    import importlib
+    import sys
+
+    (tmp_path / "search_mod.f90").write_text(SEARCH_LOOP)
+    frontend = FortranFrontend()
+    unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:search_mod")
+    facts = frontend.analyze(unit, tmp_path)
+    candidate = NumpyTranslation().apply(unit, facts, {"root": tmp_path})
+    out = tmp_path / "emitted"
+    out.mkdir()
+    for path, content in candidate.files.items():
+        (out / path.name).write_bytes(content)
+    text = (out / "search_mod_numpy.py").read_text()
+    assert text.count("else:\n") >= 2  # both loops: the index is read after each
+    sys.path.insert(0, str(out))
+    try:
+        module = importlib.import_module("search_mod_numpy")
+        import numpy as np
+
+        z = np.array([1.0, 2.0, 3.0, 4.0, 5.0], dtype=np.float32)
+        k_found, k_scan = module.first_above(5, z, np.float32(3.5))
+        assert (k_found, k_scan) == (4, 7)  # exit at z(4); 1,3,5 then one step past
+        k_found, k_scan = module.first_above(5, z, np.float32(9.0))
+        assert (k_found, k_scan) == (6, 7)  # completed: one past n
+    finally:
+        sys.path.remove(str(out))
+        sys.modules.pop("search_mod_numpy", None)

From 2ffbee5057a37824f5227fb9aa0fbdaed407cdaf Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Thu, 3 Sep 2026 21:31:17 -0600
Subject: [PATCH 07/73] SUM accumulates in element order

gfortran's inlined SUM is a loop in element order; np.sum pairs its terms
and rounds differently -- CLUBB's vertical_integral drifted 12 ULP on 3
of 7770 points. _f_vsum, beside _f_vdot, walks the array in Fortran order
or along the named axis, and the intrinsic table points SUM at it.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/numpy/runtime.py    | 17 +++++++++++++++++
 src/recast/transform/numpy/vocabulary.py |  2 +-
 tests/test_numpy_runtime.py              | 21 +++++++++++++++++++++
 3 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/src/recast/transform/numpy/runtime.py b/src/recast/transform/numpy/runtime.py
index 95fc7ce..7b73c17 100644
--- a/src/recast/transform/numpy/runtime.py
+++ b/src/recast/transform/numpy/runtime.py
@@ -203,6 +203,23 @@ def _f_vdot(a: Any, b: Any) -> Any:
     return np.dot(a, b)
 
 
+def _f_vsum(a: Any, axis: Any = None) -> Any:
+    """Fortran SUM accumulates in element order; np.sum pairs terms and
+    rounds differently (CLUBB's vertical_integral: 12 ULP)."""
+    if _LIBM_STRICT:
+        arr = np.asarray(a)
+        if axis is None:
+            s = arr.dtype.type(0) if arr.dtype.kind in "fc" else 0
+            for x in np.ravel(arr, order="F"):
+                s = s + x
+            return s
+        out = np.zeros(arr.shape[:axis] + arr.shape[axis + 1 :], dtype=arr.dtype)
+        for i in range(arr.shape[axis]):
+            out = out + np.take(arr, i, axis=axis)
+        return out
+    return np.sum(a, axis=axis)
+
+
 def _fstr_eq(a: str, b: str) -> bool:
     """Fortran character equality: pad shorter operand with blanks."""
     return a.rstrip(" ") == b.rstrip(" ")
diff --git a/src/recast/transform/numpy/vocabulary.py b/src/recast/transform/numpy/vocabulary.py
index c4441c2..7074c95 100644
--- a/src/recast/transform/numpy/vocabulary.py
+++ b/src/recast/transform/numpy/vocabulary.py
@@ -166,7 +166,7 @@
     "minval": "np.min",
     "product": "np.prod",
     "size": "np.size",
-    "sum": "np.sum",
+    "sum": "_f_vsum",
     # With unit lower bounds -- which every translated array has -- the upper
     # bound and the extent are the same number.
     "ubound": "np.size",
diff --git a/tests/test_numpy_runtime.py b/tests/test_numpy_runtime.py
index 2be1777..f068da0 100644
--- a/tests/test_numpy_runtime.py
+++ b/tests/test_numpy_runtime.py
@@ -339,3 +339,24 @@ def test_copy_out_writes_the_overlap_and_leaves_the_rest() -> None:
     runtime._f_copy_out(same, 7.0)
     assert list(same) == [7.0, 7.0]
     runtime._f_copy_out(None, np.ones(2))  # nothing to write into, no error
+
+
+def test_sum_accumulates_in_fortran_element_order() -> None:
+    """gfortran's inlined SUM is a loop in element order; np.sum pairs its
+    terms and rounds differently -- CLUBB's vertical_integral drifted 12 ULP.
+    The sequential helper matches the loop exactly, whole or along an axis."""
+    import numpy as np
+
+    from recast.transform.numpy import runtime
+
+    rng = np.random.default_rng(7)
+    a = np.asfortranarray(rng.uniform(-1e6, 1e6, size=(37, 23)))
+    loop = np.float64(0)
+    for x in np.ravel(a, order="F"):
+        loop = loop + x
+    assert runtime._f_vsum(a) == loop
+    along = np.zeros(23)
+    for i in range(37):
+        along = along + a[i, :]
+    assert np.array_equal(runtime._f_vsum(a, axis=0), along)
+    assert runtime._f_vsum(np.array([1, 2, 3], dtype=np.int32)) == 6

From 13e00d22dc57d70c671988ac17e82e9a974ab3a0 Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Thu, 3 Sep 2026 21:43:48 -0600
Subject: [PATCH 08/73] harvest's annotation says what it returns

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/fortran/use.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/recast/fortran/use.py b/src/recast/fortran/use.py
index e2424de..947a9d8 100644
--- a/src/recast/fortran/use.py
+++ b/src/recast/fortran/use.py
@@ -25,8 +25,9 @@ class UnresolvedConstant(RecastError):
     """A use-imported name whose initializer is in none of the given sources."""
 
 
-def harvest(path: Path) -> dict[str, tuple[Any, int | None]]:
-    """``name -> (initializer node, line)`` for module-level initialized entities.
+def harvest(path: Path) -> dict[str, tuple[Any, int | None, str | None]]:
+    """``name -> (initializer node, line, declared base type)`` for module-level
+    initialized entities; the base type is ``real``, ``int`` or ``None``.
 
     Covers parameters and initialized ``save``/``protected`` variables alike: a
     constant that a physics module reads is a constant whether or not the

From cd1c977b8d07d30cf922b52d9bbdcdf1ea8cf365 Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Thu, 3 Sep 2026 22:03:03 -0600
Subject: [PATCH 09/73] The completion value only where the index is read
 before it is redefined; the corpus re-baselined

An else on every loop whose variable any later statement names put one on
most loops of a module -- i and k are reused by the next loop -- and moved
the translator differential from 33 to 130 differing subprograms. A later
DO over the same variable or an assignment to it redefines it first and
closes the question; a later loop's bounds may still read it and count.
The differential is at 45, every one of the twelve new differences an
intended class: an array-valued .or. spelled |, SUM through _f_vsum, the
completion value on a search loop.

The corpus, re-run on this branch: no unit regresses; csplines, fitpack and
slsqp_core have fewer read/write disagreements; numfor's sorting reaches
the bit-exact gate for the first time (its specifics are private behind a
public generic) and exposes a bare use-imported constant in a local
parameter, issue #28.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 corpus/baseline.json                     | 54 ++++++++++++++----------
 src/recast/fortran/expr.py               |  4 +-
 src/recast/transform/numpy/statements.py | 34 +++++++++++----
 tests/test_numpy_translate.py            |  8 +++-
 4 files changed, 68 insertions(+), 32 deletions(-)

diff --git a/corpus/baseline.json b/corpus/baseline.json
index 33a3b5d..1912df3 100644
--- a/corpus/baseline.json
+++ b/corpus/baseline.json
@@ -1107,11 +1107,11 @@
      "stopped_by": "static.rwset",
      "verdicts": {
       "static.rwset": {
-       "detail": "7/66 blocks disagree: cspl_interp/B002, cspleps/B003, csplint/B006, csplint/B012, csplint/B016 (+2 more)",
+       "detail": "5/66 blocks disagree: csplint/B006, csplint/B012, csplint/B016, csplint_square/B011, csplint_square/B015",
        "metrics": {
         "blocks_checked": 66,
         "blocks_deferred": 14,
-        "blocks_matched": 59,
+        "blocks_matched": 61,
         "blocks_waived": 0
        },
        "passed": false
@@ -1158,11 +1158,11 @@
      "stopped_by": "static.rwset",
      "verdicts": {
       "static.rwset": {
-       "detail": "18/124 blocks disagree: splrep_msg/B002, splprep/B014, splprep/B016, splprep/B019, splprep/B020 (+13 more)",
+       "detail": "1/124 blocks disagree: splrep_msg/B002",
        "metrics": {
         "blocks_checked": 124,
         "blocks_deferred": 8,
-        "blocks_matched": 106,
+        "blocks_matched": 123,
         "blocks_waived": 0
        },
        "passed": false
@@ -1644,23 +1644,40 @@
       "oracle/f2py-golden": "ok",
       "store/fs-evidence": "ok",
       "transform/translate.numpy": "ok",
-      "verifier/differential.bitexact": "ok",
-      "verifier/static.rwset": "ok",
-      "verifier/symbolic.notary": "ok"
+      "verifier/differential.bitexact": "failed",
+      "verifier/static.rwset": "ok"
      },
-     "stopped_by": null,
+     "stopped_by": "differential.bitexact",
      "verdicts": {
       "differential.bitexact": {
-       "detail": "80 points across 1 subprogram(s), all bit-exact",
+       "detail": "3 subprogram(s) could not be compared: searchsorted_dp: candidate raised: NameError: name 'SMALL' is not defined; searchsorted_idp: candidate raised: NameError:",
        "metrics": {
-        "bit_exact": 80,
+        "bit_exact": 90,
         "integer_mismatch": 0,
-        "integer_points": 0,
+        "integer_points": 10,
         "max_rel": 0.0,
         "max_ulp": 0,
         "nan_mismatch": 0,
-        "points": 80,
+        "points": 90,
         "subprograms": {
+         "searchsorted_dp": {
+          "error": "candidate raised: NameError: name 'SMALL' is not defined"
+         },
+         "searchsorted_dpi": {
+          "error": "candidate raised: NameError: name 'SMALL' is not defined"
+         },
+         "searchsorted_i": {
+          "bit_exact": 10,
+          "integer_mismatch": 0,
+          "integer_points": 10,
+          "max_rel": 0.0,
+          "max_ulp": 0,
+          "nan_mismatch": 0,
+          "points": 10
+         },
+         "searchsorted_idp": {
+          "error": "candidate raised: NameError: name 'SMALL' is not defined"
+         },
          "sort": {
           "bit_exact": 80,
           "integer_mismatch": 0,
@@ -1673,7 +1690,7 @@
         },
         "trials": 10
        },
-       "passed": true
+       "passed": false
       },
       "static.rwset": {
        "detail": "62 blocks match",
@@ -1684,13 +1701,6 @@
         "blocks_waived": 0
        },
        "passed": true
-      },
-      "symbolic.notary": {
-       "detail": "no rewrites to notarize; the translation is print-order faithful",
-       "metrics": {
-        "rewrites": 0
-       },
-       "passed": true
       }
      }
     },
@@ -2098,11 +2108,11 @@
      "stopped_by": "static.rwset",
      "verdicts": {
       "static.rwset": {
-       "detail": "18/111 blocks disagree: slsqp/B019, slsqpb/B002, slsqpb/B003, slsqpb/B006, reset_bfgs_matrix/B002 (+13 more)",
+       "detail": "12/111 blocks disagree: slsqp/B019, slsqpb/B002, slsqpb/B003, slsqpb/B006, reset_bfgs_matrix/B002 (+7 more)",
        "metrics": {
         "blocks_checked": 111,
         "blocks_deferred": 3,
-        "blocks_matched": 93,
+        "blocks_matched": 99,
         "blocks_waived": 0
        },
        "passed": false
diff --git a/src/recast/fortran/expr.py b/src/recast/fortran/expr.py
index f3f3aaf..e93076b 100644
--- a/src/recast/fortran/expr.py
+++ b/src/recast/fortran/expr.py
@@ -156,7 +156,9 @@ def substitute(expr: Expr, name: str, replacement: Expr) -> Expr:
     For the one legal self-reference in an initializer, a kind inquiry on the
     constant being declared (``tol = max( 1.e-10_r8, epsilon(tol) )``): the
     reference carries the constant's kind and nothing else, and the fold
-    renders reals as 64-bit, so a 64-bit literal stands in for it.
+    renders reals as 64-bit, so a 64-bit literal stands in for it. Every
+    occurrence is replaced: a parameter cannot name itself anywhere else in
+    its own initializer, so there is no other occurrence to preserve.
     """
     if expr.kind == "name" and expr.text == name:
         return replacement
diff --git a/src/recast/transform/numpy/statements.py b/src/recast/transform/numpy/statements.py
index fbd22b6..3065c91 100644
--- a/src/recast/transform/numpy/statements.py
+++ b/src/recast/transform/numpy/statements.py
@@ -132,26 +132,44 @@ def derived_array(type_name: str, extents: list[str], known: dict[str, Any]) ->
 
 
 def _loops_whose_index_is_read_after(subprogram: Any) -> set[int]:
-    """The DO constructs whose loop variable some later statement of the
-    subprogram names -- read after the loop, where Fortran's completion
-    value (one step past the end) and Python's (the last value) differ."""
+    """The DO constructs whose loop variable a later statement reads before
+    anything redefines it -- where Fortran's completion value (one step past
+    the end) and Python's (the last value) differ and are observed.
+
+    A later DO over the same variable, or an assignment to it, redefines it
+    first and closes the question; ``i`` and ``k`` are reused by most loops
+    of a module, and an ``else`` on every one of them would be noise. The
+    search loops that read their index (CLUBB's ``lscale_width_vert_avg``)
+    are what this finds."""
     from recast.fortran.interface import names_in, node_span
 
+    def index_of(statement: Any) -> str | None:
+        control = walk(statement, f03.Loop_Control)
+        if control and control[0].children[1] is not None:
+            return str(control[0].children[1][0]).lower()
+        return None
+
     marked: set[int] = set()
     for loop in walk(subprogram, (f03.Block_Nonlabel_Do_Construct, f03.Block_Label_Do_Construct)):
         do_statement = walk(loop, (f03.Nonlabel_Do_Stmt, f03.Label_Do_Stmt))
-        control = walk(do_statement[0], f03.Loop_Control) if do_statement else []
-        if not control or control[0].children[1] is None:
-            continue
-        variable = str(control[0].children[1][0]).lower()
+        variable = index_of(do_statement[0]) if do_statement else None
         _, end_line = node_span(loop)
-        if end_line is None:
+        if variable is None or end_line is None:
             continue
         for statement in walk(subprogram):
             item = getattr(statement, "item", None)
             span = getattr(item, "span", None) if item is not None else None
             if not span or span[0] <= end_line:
                 continue
+            if isinstance(statement, (f03.Nonlabel_Do_Stmt, f03.Label_Do_Stmt)):
+                if index_of(statement) == variable:
+                    break  # redefined by the next loop over it
+                # Another loop's header may still read it in its bounds
+                # (``do k_avg = k_avg_lower, k_avg_upper``): checked below.
+            if isinstance(statement, f03.Assignment_Stmt):
+                target = statement.children[0]
+                if isinstance(target, f03.Name) and str(target).lower() == variable:
+                    break  # redefined by assignment
             if variable in names_in(statement):
                 marked.add(id(loop))
                 break
diff --git a/tests/test_numpy_translate.py b/tests/test_numpy_translate.py
index 6151e2a..22054bc 100644
--- a/tests/test_numpy_translate.py
+++ b/tests/test_numpy_translate.py
@@ -310,6 +310,10 @@ def test_logical_operators_on_arrays_are_elementwise(tmp_path: Path) -> None:
     do kk = 1, n, 2
       k_scan = kk
     end do
+    ! The bounds of a loop over another variable read kk: still a read.
+    do k = 1, kk
+      k_scan = k_scan + 0
+    end do
     k_scan = kk
   end subroutine first_above
 end module search_mod
@@ -335,7 +339,9 @@ def test_a_loop_index_read_after_the_loop_has_the_completion_value(tmp_path: Pat
     for path, content in candidate.files.items():
         (out / path.name).write_bytes(content)
     text = (out / "search_mod_numpy.py").read_text()
-    assert text.count("else:\n") >= 2  # both loops: the index is read after each
+    # The first two loops' indices are read after them (one in a later
+    # loop's bounds); the last loop's index k is not.
+    assert text.count("max(0, ") == 2
     sys.path.insert(0, str(out))
     try:
         module = importlib.import_module("search_mod_numpy")

From 63c003615367b97822c46621d588817086fcdf5d Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Thu, 3 Sep 2026 22:47:26 -0600
Subject: [PATCH 10/73] What the tier-2 plans needed: uncarried strings, listed
 state, extents the run owns

A character component (err_info%err_header) has no flat spelling and the
physics reads it only to print: left at the object's default, under
left_to_module. A derived-type module variable declared in a list over
several lines (the four sponge settings) is found on the module record
when the one-line pattern misses it. An allocation sized by the allocating
routine's dummy when the planned subprogram has none of that name
(coef_wp4_implicit(1:ngrdcol,1:nz) under advance_wp2_wp3, which takes
nzm and nzt) becomes an integer argument of the adapter, extent_args,
which the recorder writes from size() and the sampled gate sizes like any
scalar an array's dims name. Every advance_* plan of CLUBB is usable on
this.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/fortran/flatten.py | 58 ++++++++++++++++++++++++++++----
 src/recast/oracle/record.py   |  5 +++
 tests/test_flatten.py         | 63 +++++++++++++++++++++++++++++++++++
 3 files changed, 119 insertions(+), 7 deletions(-)

diff --git a/src/recast/fortran/flatten.py b/src/recast/fortran/flatten.py
index d3dde44..9adfb48 100644
--- a/src/recast/fortran/flatten.py
+++ b/src/recast/fortran/flatten.py
@@ -146,11 +146,20 @@ class FlatPlan:
     dim_constants: dict[str, int] = field(default_factory=dict)
     """Named extents of the original dummies that are tree constants
     (``a(nrk,nrk)``), so the flat signature can spell them as numbers."""
+    extent_args: dict[str, list[Any]] = field(default_factory=dict)
+    """``{name: [object, component, axis]}``: an integer argument of the
+    adapter carrying an extent the plan could not spell -- an allocation
+    sized by the allocating routine's dummy when the planned subprogram
+    has none of that name (CLUBB's ``coef_wp4_implicit(1:ngrdcol, 1:nz)``
+    under ``advance_wp2_wp3``, which takes ``nzm`` and ``nzt``). The
+    recorder writes it from ``size()``; the sampled gate sizes it like any
+    scalar an array's dims name."""
+
     left_to_module: list[str] = field(default_factory=list)
-    """Module state the body reaches that the adapter cannot set: a
-    variable its module keeps private (CLUBB's ``error_code %
-    clubb_debug_level``). Both sides run with the module's own default,
-    and the plan says so rather than failing."""
+    """What the body reaches that the adapter does not carry: module state
+    its module keeps private (CLUBB's ``error_code % clubb_debug_level``),
+    a character component (``err_info % err_header``). Both sides run with
+    the default, and the plan says so rather than failing."""
     patch_count: str = "np_"
     counter_prefix: str = "num_"
 
@@ -201,6 +210,7 @@ def from_dict(cls, data: dict[str, Any]) -> FlatPlan:
             states=states,
             dim_constants=dict(data.get("dim_constants", {})),
             left_to_module=list(data.get("left_to_module", [])),
+            extent_args={k: list(v) for k, v in (data.get("extent_args") or {}).items()},
             patch_count=data.get("patch_count", "np_"),
             counter_prefix=data.get("counter_prefix", "num_"),
         )
@@ -252,6 +262,10 @@ def flat_args(self) -> list[dict[str, Any]]:
                     "dims": None,
                 }
             )
+        for name in self.extent_args:
+            args.append(
+                {"name": name, "dtype": "int32", "intent": "IN", "optional": False, "dims": None}
+            )
         for obj in self.objects:
             for comp in obj.components:
                 args.append(
@@ -338,6 +352,17 @@ def _state_declaration(name: str, root: Path) -> tuple[str, str] | None:
                 continue  # a dummy of that type, not the module's variable
             module = MODULE_DEFINITION.search(text)
             return match.group(1).lower(), (module.group(1).lower() if module else "")
+    # Declared in a list over several lines (CLUBB's four sponge settings
+    # under one ``type(sponge_damp_settings), public :: &``): the module
+    # record has read the declaration whole.
+    for path in sources(root):
+        record = _module_record(path, {})
+        for entry in (record or {}).get("module_state", ()):
+            if str(entry.get("name", "")).lower() != name.lower():
+                continue
+            derived = DERIVED.match(str(entry.get("dtype", "")))
+            if derived:
+                return derived.group(1).lower(), str(record.get("module", "")).lower()
     return None
 
 
@@ -919,6 +944,12 @@ def type_info(type_name: str) -> tuple[dict[str, Any], dict[str, list[str]]] | N
                 if spec is None:
                     plan.unsupported.append(f"{obj}%{member}: no such component")
                     continue
+                if str(spec.get("dtype")) not in FORTRAN_TYPES:
+                    # A character component (CLUBB's err_info%err_header): the
+                    # adapter has no flat spelling for it and the physics reads
+                    # it only to print. Left at the object's default, and said.
+                    plan.left_to_module.append(f"{obj}%{member}: {spec.get('dtype')} not carried")
+                    continue
                 found_axes = bounds.get(member)
                 if found_axes is None and spec.get("dims"):
                     plan.unsupported.append(f"{obj}%{member}: no allocate statement found")
@@ -1104,7 +1135,9 @@ def _bind_symbolic_extents(plan: FlatPlan) -> None:
     carried = {(obj.name, comp.name): comp.flat for obj in plan.objects for comp in obj.components}
     flat_names = set(carried.values())
 
-    def bind(text: str) -> str | None:
+    def bind(text: str, synthetic: str | None = None) -> str | None:
+        """``synthetic`` names the extent argument to stand in when the
+        text is one identifier nothing else answers for."""
         missing: list[str] = []
 
         def component(match: re.Match[str]) -> str:
@@ -1129,13 +1162,24 @@ def swap(match: re.Match[str]) -> str:
             return token
 
         out = re.sub(r"[A-Za-z_]\w*", swap, text)
+        if missing and synthetic is not None:
+            bare = text.strip().strip("()").strip()
+            if re.fullmatch(r"[A-Za-z_]\w*", bare) and bare.lower() == missing[0].lower():
+                return synthetic
         return None if missing else out
 
     for obj in plan.objects:
         kept = []
         for comp in obj.components:
-            extents = [bind(e) for e in comp.extents]
-            bounds = [(bind(lo), bind(hi)) for lo, hi in comp.bounds]
+            names = [f"{comp.flat}_n{axis + 1}" for axis in range(len(comp.extents))]
+            extents = [bind(e, names[axis]) for axis, e in enumerate(comp.extents)]
+            bounds = [
+                (bind(lo), bind(hi, names[axis]) if lo.strip() == "1" else bind(hi))
+                for axis, (lo, hi) in enumerate(comp.bounds)
+            ]
+            for axis, extent in enumerate(extents):
+                if extent == names[axis]:
+                    plan.extent_args[names[axis]] = [obj.name, comp.name, axis + 1]
             if any(e is None for e in extents) or any(
                 lo is None or hi is None for lo, hi in bounds
             ):
diff --git a/src/recast/oracle/record.py b/src/recast/oracle/record.py
index 808022e..acc174e 100644
--- a/src/recast/oracle/record.py
+++ b/src/recast/oracle/record.py
@@ -244,6 +244,11 @@ def recorder_module(
             fmt = {"int32": "i0", "bool": "l1"}.get(str(a["dtype"]), "es25.17e3")
             lines.append(f"       write ({u}, '(a,{fmt})') '# {a['name']} = ', {a['name']}")
         lines.append(f"       write ({u}, '(a,i0)') '# {patch} = ', {patch}")
+        for extent, (owner, member, axis) in plan.extent_args.items():
+            # An extent the plan could not spell: the run's own, from size().
+            lines.append(
+                f"       write ({u}, '(a,i0)') '# {extent} = ', size({owner}%{member}, {axis})"
+            )
 
         for a in originals:
             if a.get("dims") and not DERIVED.match(str(a["dtype"])):
diff --git a/tests/test_flatten.py b/tests/test_flatten.py
index c68ba1f..f14f0ae 100644
--- a/tests/test_flatten.py
+++ b/tests/test_flatten.py
@@ -567,3 +567,66 @@ def test_an_object_allocated_many_at_once_and_sized_by_itself(tmp_path: Path) ->
     text = fortran_adapter("column_mod", [plan], [])
     assert "real(8), intent(in) :: gr__zm(ngrdcol, gr__nzm)" in text
     assert "allocate(gr%zm(1:ngrdcol, 1:gr__nzm))" in text
+
+
+COEFS = """\
+module coefs_mod
+  implicit none
+  private
+  public :: coefs_type, init_coefs
+  type coefs_type
+    real(8), allocatable, dimension(:,:) :: coef
+  end type coefs_type
+contains
+  subroutine init_coefs( ngrdcol, nz, c )
+    integer, intent(in) :: ngrdcol, nz
+    type(coefs_type), intent(out) :: c
+    allocate( c%coef(1:ngrdcol,1:nz) )
+    c%coef = 0.0d0
+  end subroutine init_coefs
+end module coefs_mod
+"""
+
+USES_COEFS = """\
+module solver_mod
+  use coefs_mod, only: coefs_type
+  implicit none
+  private
+  public :: apply
+contains
+  subroutine apply( nzt, ngrdcol, c, x )
+    integer, intent(in) :: nzt, ngrdcol
+    type(coefs_type), intent(in) :: c
+    real(8), intent(inout), dimension(ngrdcol, nzt) :: x
+    x = x * c%coef(:, 1:nzt)
+  end subroutine apply
+end module solver_mod
+"""
+
+
+def test_an_extent_the_plan_cannot_spell_becomes_an_argument(tmp_path: Path) -> None:
+    """``coef`` is allocated ``(1:ngrdcol, 1:nz)`` by the initializer's dummy
+    ``nz``; the planned subprogram takes ``nzt``, not ``nz``. The extent is
+    the run's own: the plan makes it an integer argument, the recorder writes
+    it from ``size()``, and the adapters declare the component by it."""
+    from recast.oracle.record import recorder_module
+
+    (tmp_path / "coefs_mod.f90").write_text(COEFS)
+    (tmp_path / "solver_mod.f90").write_text(USES_COEFS)
+    frontend = FortranFrontend(flatten={"patch_count": "ngrdcol", "bounds_pattern": r"^ngrdcol$"})
+    unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:solver_mod")
+    facts = frontend.analyze(unit, tmp_path)
+    (plan,) = plans_for(facts, tmp_path, CLUBB_CONVENTIONS)
+    assert plan.usable, plan.unsupported
+    assert plan.extent_args == {"c__coef_n2": ["c", "coef", 2]}
+    (coef,) = plan.objects[0].components
+    assert coef.extents == ["ngrdcol", "c__coef_n2"]
+    names = [a["name"] for a in plan.flat_args]
+    assert names.index("c__coef_n2") < names.index("c__coef")
+    adapter = fortran_adapter("solver_mod", [plan], [])
+    assert "integer, intent(in) :: c__coef_n2" in adapter
+    assert "real(8), intent(in) :: c__coef(ngrdcol, c__coef_n2)" in adapter
+    recorder = recorder_module("solver_mod", [plan])
+    assert "'# c__coef_n2 = ', size(c%coef, 2)" in recorder
+    again = FlatPlan.from_dict(plan.to_dict())
+    assert again.extent_args == plan.extent_args

From 473afa32551146241d769aac2f81f8b7b34d4f5a Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Thu, 3 Sep 2026 22:54:42 -0600
Subject: [PATCH 11/73] The recorder declares the patch count once when the
 probe already takes it

CLUBB passes ngrdcol as a dummy of every subroutine the recorder probes;
a local of the same name was a second declaration, and gfortran said so
forty times.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/oracle/record.py | 11 +++++++++--
 tests/test_flatten.py       |  4 ++++
 2 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/src/recast/oracle/record.py b/src/recast/oracle/record.py
index acc174e..cfa41b7 100644
--- a/src/recast/oracle/record.py
+++ b/src/recast/oracle/record.py
@@ -211,7 +211,12 @@ def recorder_module(
                 spelled = FORTRAN_TYPES[str(a["dtype"])]
                 dims = "(" + ",".join(":" for _ in a["dims"]) + ")" if a.get("dims") else ""
                 lines.append(f"    {spelled}, intent(in) :: {a['name']}{dims}")
-        lines.append(f"    integer :: {patch}")
+        passes_patch = any(a["name"].lower() == patch.lower() for a in originals)
+        if not passes_patch:
+            # CLUBB passes ngrdcol as a dummy; then it is already declared and
+            # already the run's value, and a local of the same name would be
+            # a second declaration.
+            lines.append(f"    integer :: {patch}")
         lines.append("    character(len=128) :: dims")
         # The patch count from the first component allocated over it.
         first = next(
@@ -223,7 +228,9 @@ def recorder_module(
             ),
             None,
         )
-        if first is None:
+        if passes_patch:
+            pass
+        elif first is None:
             lines.append(f"    {patch} = 1")
         else:
             lines.append(f"    {patch} = size({first[0].name}%{first[1].name}, 1)")
diff --git a/tests/test_flatten.py b/tests/test_flatten.py
index f14f0ae..0e54540 100644
--- a/tests/test_flatten.py
+++ b/tests/test_flatten.py
@@ -628,5 +628,9 @@ def test_an_extent_the_plan_cannot_spell_becomes_an_argument(tmp_path: Path) ->
     assert "real(8), intent(in) :: c__coef(ngrdcol, c__coef_n2)" in adapter
     recorder = recorder_module("solver_mod", [plan])
     assert "'# c__coef_n2 = ', size(c%coef, 2)" in recorder
+    # ngrdcol is a dummy of the probe already: declared once, not assigned.
+    probe = recorder[recorder.index("subroutine rec_apply(") :]
+    assert probe.count("integer :: ngrdcol") == 0
+    assert "ngrdcol = " not in probe.split("phase == 0")[0]
     again = FlatPlan.from_dict(plan.to_dict())
     assert again.extent_args == plan.extent_args

From a11485951e22b0efea96929b370de7aa6b7da6f5 Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Thu, 3 Sep 2026 22:58:20 -0600
Subject: [PATCH 12/73] A call continued with trailing comments is probed whole

CLUBB continues its calls as a, b, & ! In on every line: the comment after
the ampersand left a blank the joiner did not strip, and the probe carried
& into its argument list.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/oracle/record.py |  6 +++++-
 tests/test_flatten.py       | 37 +++++++++++++++++++++++++++++++++++++
 2 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/src/recast/oracle/record.py b/src/recast/oracle/record.py
index cfa41b7..0d09db7 100644
--- a/src/recast/oracle/record.py
+++ b/src/recast/oracle/record.py
@@ -389,7 +389,11 @@ def probe_tree(
             span = [lines[i]]
             while _continues(span[-1]) and i + len(span) < len(lines):
                 span.append(lines[i + len(span)])
-            logical = " ".join(_strip_comment(ln).rstrip("&").strip().lstrip("&") for ln in span)
+            # ``a, & ! In`` leaves a space between the ampersand and the
+            # comment it trailed: strip blanks before the ampersands.
+            logical = " ".join(
+                _strip_comment(ln).strip().rstrip("&").strip().lstrip("&").strip() for ln in span
+            )
             match = CALL.match(" " * (len(span[0]) - len(span[0].lstrip())) + logical.strip())
             called = match.group("name").lower() if match else None
             if match and called and called in targets:
diff --git a/tests/test_flatten.py b/tests/test_flatten.py
index 0e54540..06572e7 100644
--- a/tests/test_flatten.py
+++ b/tests/test_flatten.py
@@ -634,3 +634,40 @@ def test_an_extent_the_plan_cannot_spell_becomes_an_argument(tmp_path: Path) ->
     assert "ngrdcol = " not in probe.split("phase == 0")[0]
     again = FlatPlan.from_dict(plan.to_dict())
     assert again.extent_args == plan.extent_args
+
+
+def test_a_call_continued_with_trailing_comments_is_probed_whole(tmp_path: Path) -> None:
+    """CLUBB continues its calls as ``a, b, & ! In`` on every line: the
+    comment after the ampersand left a blank the joiner did not strip, and
+    the probe carried ``&`` into its argument list."""
+    from recast.oracle.record import probe_tree
+
+    (tmp_path / "coefs_mod.f90").write_text(COEFS)
+    (tmp_path / "solver_mod.f90").write_text(USES_COEFS)
+    (tmp_path / "driver_mod.f90").write_text(
+        """\
+module driver_mod
+  use coefs_mod, only: coefs_type
+  use solver_mod, only: apply
+  implicit none
+contains
+  subroutine step( nzt, ngrdcol, c, x )
+    integer, intent(in) :: nzt, ngrdcol
+    type(coefs_type), intent(in) :: c
+    real(8), intent(inout), dimension(ngrdcol, nzt) :: x
+    call apply( nzt, ngrdcol, & ! In
+                c, &        ! In
+                x )         ! In/out
+  end subroutine step
+end module driver_mod
+"""
+    )
+    frontend = FortranFrontend(flatten={"patch_count": "ngrdcol", "bounds_pattern": r"^ngrdcol$"})
+    unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:solver_mod")
+    facts = frontend.analyze(unit, tmp_path)
+    plans = plans_for(facts, tmp_path, CLUBB_CONVENTIONS)
+    sites = probe_tree(tmp_path, tmp_path / "probed", {"solver_mod": plans})
+    assert sites == {"apply": 1}
+    probed = (tmp_path / "probed" / "driver_mod.f90").read_text()
+    assert "call rec_apply(0, nzt, ngrdcol, c, x)" in probed
+    assert "&" not in probed.split("call rec_apply(0")[1].split("\n")[0]

From a67d712e4396cfaa16726e7875e5c90070c3ba75 Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 11:04:44 -0600
Subject: [PATCH 13/73] The recorder guards a component the run may not
 allocate

CLUBB allocates its scalar-tracer coefficients only when sclr_dim > 0;
under bomex it is 0, and reshape of an unallocated component faulted the
probed run after three dumps. An allocatable or pointer component is
recorded inside allocated()/associated(), as a zero-extent record
otherwise, and the extent argument taken from its size() is zero when it
was never allocated.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/oracle/record.py | 54 ++++++++++++++++++++-----------------
 tests/test_flatten.py       | 19 ++++++++++++-
 2 files changed, 47 insertions(+), 26 deletions(-)

diff --git a/src/recast/oracle/record.py b/src/recast/oracle/record.py
index 0d09db7..535e1de 100644
--- a/src/recast/oracle/record.py
+++ b/src/recast/oracle/record.py
@@ -96,6 +96,25 @@ def _record_call(
     )
 
 
+def _record_component(lines: list[str], unit: str, tag: str, obj: Any, comp: Any) -> None:
+    """One component's record, guarded when it may not be there: a component
+    the model allocates only on some configurations (CLUBB's scalar-tracer
+    coefficients under ``sclr_dim = 0``) is written as a zero-extent record,
+    and ``reshape`` is never asked for storage that does not exist."""
+    target = f"{obj.name}%{comp.name}"
+    rank = len(comp.extents)
+    if not comp.bounds:
+        _record_call(lines, unit, tag, comp.flat, target, rank, comp.dtype)
+        return
+    test = "associated" if comp.pointer else "allocated"
+    lines.append(f"       if ({test}({target})) then")
+    _record_call(lines, unit, tag, comp.flat, target, rank, comp.dtype)
+    lines.append("       else")
+    zeros = "(" + ",".join("0" for _ in range(rank)) + ")" if rank else ""
+    lines.append(f"       write ({unit}, '(a)') '# {tag}: {comp.flat}{zeros}'")
+    lines.append("       end if")
+
+
 def recorder_module(
     module: str,
     plans: list[FlatPlan],
@@ -252,9 +271,15 @@ def recorder_module(
             lines.append(f"       write ({u}, '(a,{fmt})') '# {a['name']} = ', {a['name']}")
         lines.append(f"       write ({u}, '(a,i0)') '# {patch} = ', {patch}")
         for extent, (owner, member, axis) in plan.extent_args.items():
-            # An extent the plan could not spell: the run's own, from size().
+            # An extent the plan could not spell: the run's own, from size(),
+            # zero when the component was never allocated.
+            component = next(
+                c for o in plan.objects if o.name == owner for c in o.components if c.name == member
+            )
+            test = "associated" if component.pointer else "allocated"
             lines.append(
-                f"       write ({u}, '(a,i0)') '# {extent} = ', size({owner}%{member}, {axis})"
+                f"       write ({u}, '(a,i0)') '# {extent} = ', "
+                f"merge(size({owner}%{member}, {axis}), 0, {test}({owner}%{member}))"
             )
 
         for a in originals:
@@ -264,15 +289,7 @@ def recorder_module(
                 )
         for obj in plan.objects:
             for comp in obj.components:
-                _record_call(
-                    lines,
-                    u,
-                    "INPUT",
-                    comp.flat,
-                    f"{obj.name}%{comp.name}",
-                    len(comp.extents),
-                    comp.dtype,
-                )
+                _record_component(lines, u, "INPUT", obj, comp)
         for state in plan.states:
             _record_call(
                 lines,
@@ -288,20 +305,7 @@ def recorder_module(
         for obj in plan.objects:
             for comp in obj.components:
                 if comp.written:
-                    _record_call(
-                        lines,
-                        u,
-                        "OUTPUT",
-                        comp.flat,
-                        f"{obj.name}%{comp.name}",
-                        len(comp.extents),
-                        comp.dtype,
-                    )
-        for a in originals:
-            if a.get("dims") and not DERIVED.match(str(a["dtype"])) and a["intent"] != "IN":
-                _record_call(
-                    lines, u, "OUTPUT", a["name"], a["name"], len(a["dims"]), str(a["dtype"])
-                )
+                    _record_component(lines, u, "OUTPUT", obj, comp)
         for state in plan.states:
             if state.written:
                 _record_call(
diff --git a/tests/test_flatten.py b/tests/test_flatten.py
index 06572e7..c557a27 100644
--- a/tests/test_flatten.py
+++ b/tests/test_flatten.py
@@ -627,7 +627,7 @@ def test_an_extent_the_plan_cannot_spell_becomes_an_argument(tmp_path: Path) ->
     assert "integer, intent(in) :: c__coef_n2" in adapter
     assert "real(8), intent(in) :: c__coef(ngrdcol, c__coef_n2)" in adapter
     recorder = recorder_module("solver_mod", [plan])
-    assert "'# c__coef_n2 = ', size(c%coef, 2)" in recorder
+    assert "'# c__coef_n2 = ', merge(size(c%coef, 2), 0, allocated(c%coef))" in recorder
     # ngrdcol is a dummy of the probe already: declared once, not assigned.
     probe = recorder[recorder.index("subroutine rec_apply(") :]
     assert probe.count("integer :: ngrdcol") == 0
@@ -671,3 +671,20 @@ def test_a_call_continued_with_trailing_comments_is_probed_whole(tmp_path: Path)
     probed = (tmp_path / "probed" / "driver_mod.f90").read_text()
     assert "call rec_apply(0, nzt, ngrdcol, c, x)" in probed
     assert "&" not in probed.split("call rec_apply(0")[1].split("\n")[0]
+
+
+def test_the_recorder_guards_a_component_the_run_may_not_allocate(tmp_path: Path) -> None:
+    """CLUBB allocates its scalar-tracer coefficients only when sclr_dim > 0;
+    reshape of an unallocated component faulted the recording run."""
+    from recast.oracle.record import recorder_module
+
+    (tmp_path / "coefs_mod.f90").write_text(COEFS)
+    (tmp_path / "solver_mod.f90").write_text(USES_COEFS)
+    frontend = FortranFrontend(flatten={"patch_count": "ngrdcol", "bounds_pattern": r"^ngrdcol$"})
+    unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:solver_mod")
+    facts = frontend.analyze(unit, tmp_path)
+    (plan,) = plans_for(facts, tmp_path, CLUBB_CONVENTIONS)
+    recorder = recorder_module("solver_mod", [plan])
+    assert "if (allocated(c%coef)) then" in recorder
+    assert "'# INPUT: c__coef(0,0)'" in recorder
+    assert "merge(size(c%coef, 2), 0, allocated(c%coef))" in recorder

From 9d482bfcce7adcf7a72bb4669844d58bed046aad Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 11:23:41 -0600
Subject: [PATCH 14/73] A sibling's generic resolves its positions by arity; a
 stubbed call reads and writes nothing

The read/write scope marked whatever sat at a union of positions as
written when a call reached a sibling's generic whose specifics differ in
arity: CLUBB's tridiag_solve and zm2zt_api marked nzt, l_implemented and
gr. The companion table now carries each specific with its argument
count, and the scope picks the one the actuals fit.

A name use-imported from a stubbed module is a call the translation
answers with a stub -- pass, for stats_update -- so the source side no
longer counts its actuals: stats_tmp and the budget names were reads only
that side saw.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/fortran/frontend.py  | 16 +++++++++++-
 src/recast/fortran/interface.py | 26 ++++++++++++++++---
 src/recast/fortran/rwset.py     |  8 ++++++
 tests/test_fortran_analysis.py  | 46 ++++++++++++++++++++++++++++++++-
 4 files changed, 90 insertions(+), 6 deletions(-)

diff --git a/src/recast/fortran/frontend.py b/src/recast/fortran/frontend.py
index c42652d..4918de1 100644
--- a/src/recast/fortran/frontend.py
+++ b/src/recast/fortran/frontend.py
@@ -394,7 +394,12 @@ def analyze(self, unit: Unit, root: Path) -> Facts:
         from recast.fortran._parse import STD, digest, f03
         from recast.fortran._parse import parse as parse_file
         from recast.fortran.effects import side_channels
-        from recast.fortran.interface import _scope_of, companion_externals, subprogram_key
+        from recast.fortran.interface import (
+            _only_names,
+            _scope_of,
+            companion_externals,
+            subprogram_key,
+        )
         from recast.fortran.rwset import block_rwsets, scope_for
 
         path = self._source_of(unit, Path(root))
@@ -475,6 +480,15 @@ def analyze(self, unit: Unit, root: Path) -> Facts:
                     table[local] = table[remote]
             for name, entry in table.items():
                 externals.setdefault(name, entry)
+        # A name use-imported from a stubbed module is a call the translation
+        # answers with a stub -- ``pass`` for CLUBB's stats_update -- so on
+        # this side too its actuals are neither read nor written.
+        for statement in record.get("use_statements", ()):
+            match = USE_STATEMENT.match(statement.strip())
+            if match and match.group("module").lower() in self.stub_modules:
+                for name in _only_names(statement):
+                    stubbed = {"kind": "subroutine", "out_positions": [], "buffer_positions": []}
+                    externals.setdefault(name, {**stubbed, "stub": True})
 
         callgraph: dict[str, list[str]] = {}
         effects: dict[str, Any] = {}
diff --git a/src/recast/fortran/interface.py b/src/recast/fortran/interface.py
index ef91ae6..be37940 100644
--- a/src/recast/fortran/interface.py
+++ b/src/recast/fortran/interface.py
@@ -1317,15 +1317,33 @@ def companion_externals(record: dict[str, Any]) -> dict[str, dict[str, Any]]:
     # writes are the union over the specifics -- which agree, in every
     # generic seen so far, on being functions with no OUT argument.
     for generic, specifics in (record.get("generics") or {}).items():
-        known = [table[s] for s in specifics if s in table]
+        known = [(s, table[s]) for s in specifics if s in table]
         if generic in table or not known:
             continue
+        # Which specific a call reaches depends on its arity (and ranks the
+        # scope does not resolve): the entry carries every specific with its
+        # argument count, and the scope picks by the actuals it sees. A
+        # union over specifics of different arity marked the wrong
+        # positions -- CLUBB's tridiag_solve, zm2zt_api.
+        arity = {
+            s: len(next(x for x in record["subprograms"] if x["name"] == s)["args"])
+            for s, _ in known
+        }
         table[generic] = {
-            "kind": known[0]["kind"],
-            "out_positions": sorted({at for entry in known for at in entry["out_positions"]}),
+            "kind": known[0][1]["kind"],
+            "out_positions": sorted({at for _, entry in known for at in entry["out_positions"]}),
             "buffer_positions": sorted(
-                {at for entry in known for at in entry.get("buffer_positions", [])}
+                {at for _, entry in known for at in entry.get("buffer_positions", [])}
             ),
+            "specifics": [
+                {
+                    "name": s,
+                    "args": arity[s],
+                    "out_positions": entry["out_positions"],
+                    "buffer_positions": entry.get("buffer_positions", []),
+                }
+                for s, entry in known
+            ],
         }
     return table
 
diff --git a/src/recast/fortran/rwset.py b/src/recast/fortran/rwset.py
index 7861b2b..377ee09 100644
--- a/src/recast/fortran/rwset.py
+++ b/src/recast/fortran/rwset.py
@@ -383,6 +383,14 @@ def call(stmt: Any) -> None:
 
         if callee is None:
             external = scope.externals.get(name)
+            if external and external.get("stub"):
+                return  # a stubbed call: the translation drops it, actuals and all
+            if external and external.get("specifics"):
+                # A sibling's generic: the specific this arity reaches, or the
+                # union when none matches exactly.
+                fitting = [x for x in external["specifics"] if x["args"] == len(actuals)]
+                if fitting:
+                    external = fitting[0]
             out_positions = set(external.get("out_positions", [])) if external else set()
             buffers = set(external.get("buffer_positions", [])) if external else set()
             for j, actual in enumerate(actuals):
diff --git a/tests/test_fortran_analysis.py b/tests/test_fortran_analysis.py
index b0ab4f6..8907d88 100644
--- a/tests/test_fortran_analysis.py
+++ b/tests/test_fortran_analysis.py
@@ -1819,7 +1819,13 @@ def test_a_siblings_generic_is_a_procedure_to_the_read_write_scope(tmp_path: Pat
     without an entry for it the scope counted the name as a read of data."""
     record = interface.extract(_write(tmp_path, "solve.f90", PUBLIC_GENERIC))
     table = interface.companion_externals(record)
-    assert table["solve"] == {"kind": "subroutine", "out_positions": [1, 2], "buffer_positions": []}
+    assert table["solve"]["kind"] == "subroutine"
+    assert table["solve"]["out_positions"] == [1, 2]  # the union, for a call of no known arity
+    # ... and each specific with its arity, for the scope to pick by the actuals.
+    assert [(x["name"], x["args"], x["out_positions"]) for x in table["solve"]["specifics"]] == [
+        ("solve_one", 2, [1]),
+        ("solve_many", 3, [2]),
+    ]
     assert table["solve_one"]["out_positions"] == [1]
 
 
@@ -1850,3 +1856,41 @@ def test_a_quotient_of_real_parameters_is_a_real_quotient(tmp_path: Path) -> Non
     assert scope["EP"] == np.float64(287.04) / np.float64(461.5)
     assert scope["EP2"] == np.float64(1.0) / scope["EP"]
     assert scope["NRK"] == 4
+
+
+STUBBED_CALLER = """\
+module budget_mod
+  use stats_mod, only: stats_type, stats_update
+  implicit none
+  private
+  public :: tend
+contains
+  subroutine tend( n, x, stats )
+    integer, intent(in) :: n
+    real, dimension(n), intent(inout) :: x
+    type(stats_type), intent(inout) :: stats
+    real, dimension(n) :: stats_tmp
+    x = 2.0 * x
+    if ( stats%l_sample ) then
+      stats_tmp = x / 2.0
+      call stats_update( "x_budget", stats_tmp, stats )
+    end if
+  end subroutine tend
+end module budget_mod
+"""
+
+
+def test_a_call_into_a_stubbed_module_reads_and_writes_nothing(tmp_path: Path) -> None:
+    """CLUBB brackets its budgets with calls into stats_netcdf, a stub: the
+    translation emits ``pass`` for them, so the source side must not count
+    their actuals either -- ``stats_tmp`` was a read only the source saw."""
+    from recast.fortran.frontend import FortranFrontend
+
+    _write(tmp_path, "budget.f90", STUBBED_CALLER)
+    frontend = FortranFrontend(stub_modules=["stats_mod"])
+    unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:budget_mod")
+    facts = frontend.analyze(unit, tmp_path)
+    assert facts.interface  # analysed without the stub module in the tree
+    blocks = facts.effects["fortran:budget_mod/tend"]["blocks"]
+    reads = {name for block in blocks for name in block.get("reads", [])}
+    assert "stats_tmp" not in reads and "x_budget" not in reads

From 34568fbac7831e4bc7c0593f753d4d597e60e8e4 Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 11:32:46 -0600
Subject: [PATCH 15/73] The stub rule covers the stub module's procedures, on
 both sides

The first cut put every name imported from a stubbed module into the
externals table, and a bare name in that table is not a read on the source
side: every constant of constants_clubb -- a stub and a table of constants
-- became a read only the target saw. The frontend now asks the stub
module's own record which imports are procedures (an absent stub is taken
at its import list), and hands the list to the translation, whose protocol
lists them as procedures so a call to a stand-in function is a call and not
a read of its name.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/fortran/frontend.py          | 52 ++++++++++++++++++++-----
 src/recast/transform/numpy/translate.py |  3 ++
 2 files changed, 45 insertions(+), 10 deletions(-)

diff --git a/src/recast/fortran/frontend.py b/src/recast/fortran/frontend.py
index 4918de1..957ab41 100644
--- a/src/recast/fortran/frontend.py
+++ b/src/recast/fortran/frontend.py
@@ -395,7 +395,6 @@ def analyze(self, unit: Unit, root: Path) -> Facts:
         from recast.fortran._parse import parse as parse_file
         from recast.fortran.effects import side_channels
         from recast.fortran.interface import (
-            _only_names,
             _scope_of,
             companion_externals,
             subprogram_key,
@@ -480,15 +479,18 @@ def analyze(self, unit: Unit, root: Path) -> Facts:
                     table[local] = table[remote]
             for name, entry in table.items():
                 externals.setdefault(name, entry)
-        # A name use-imported from a stubbed module is a call the translation
-        # answers with a stub -- ``pass`` for CLUBB's stats_update -- so on
-        # this side too its actuals are neither read nor written.
-        for statement in record.get("use_statements", ()):
-            match = USE_STATEMENT.match(statement.strip())
-            if match and match.group("module").lower() in self.stub_modules:
-                for name in _only_names(statement):
-                    stubbed = {"kind": "subroutine", "out_positions": [], "buffer_positions": []}
-                    externals.setdefault(name, {**stubbed, "stub": True})
+        # A procedure use-imported from a stubbed module is a call the
+        # translation answers with a stub -- ``pass`` for CLUBB's stats_update
+        # -- so on this side too its actuals are neither read nor written.
+        # Procedures only: a constant the same module exports (constants_clubb
+        # is a stub and a table of constants) is a read on both sides. The
+        # stub module's own record says which names are procedures; a stub
+        # the tree does not carry is taken at its import list.
+        stub_procedures = self._stub_procedures(record, Path(root))
+        for name in stub_procedures:
+            stubbed = {"kind": "subroutine", "out_positions": [], "buffer_positions": []}
+            externals.setdefault(name, {**stubbed, "stub": True})
+        record = {**record, "stub_procedures": sorted(stub_procedures)}
 
         callgraph: dict[str, list[str]] = {}
         effects: dict[str, Any] = {}
@@ -644,6 +646,36 @@ def _tree_kinds(
                 pending.extend(record_of.get("use_statements", ()))
         return found
 
+    def _stub_procedures(self, record: dict[str, Any], root: Path) -> set[str]:
+        """The local names this unit imports from stubbed modules that are
+        procedures of theirs -- calls the translation stubs. A stub module the
+        tree does not carry contributes every name it is imported for."""
+        from recast.fortran import interface as interface_mod
+
+        index = self._module_index(root.resolve())
+        names: set[str] = set()
+        for statement in record.get("use_statements", ()):
+            match = USE_STATEMENT.match(statement.strip())
+            if not match or match.group("module").lower() not in self.stub_modules:
+                continue
+            module = match.group("module").lower()
+            imported = {
+                item.split("=>", 1)[0].strip().lower(): item.split("=>", 1)[-1].strip().lower()
+                for item in (match.group("only") or "").split(",")
+                if item.strip()
+            }
+            source = index.get(module)
+            procedures = None
+            if source is not None:
+                record_of = self._readable(source, interface_mod.extract, module)
+                if record_of is not None:
+                    procedures = {str(sub["name"]).lower() for sub in record_of["subprograms"]}
+                    procedures |= {g.lower() for g in record_of.get("generics") or {}}
+            for local, remote in imported.items():
+                if procedures is None or remote in procedures:
+                    names.add(local)
+        return names
+
     def _companions(
         self, record: dict[str, Any], path: Path, root: Path
     ) -> tuple[list[dict[str, Any]], list[dict[str, str]]]:
diff --git a/src/recast/transform/numpy/translate.py b/src/recast/transform/numpy/translate.py
index 3891414..368e6d6 100644
--- a/src/recast/transform/numpy/translate.py
+++ b/src/recast/transform/numpy/translate.py
@@ -422,6 +422,9 @@ def _rwset_protocol(
                 # The siblings' procedures too: `_wv.wv_sat_svp_water(t)` is a
                 # call, and without these the alias rule would read it as data.
                 | {remote.name for remote in assembler.remotes.values()}
+                # ... and the stubbed modules' (the frontend's list): a call to
+                # a stand-in function is a call, not a read of its name.
+                | {pysafe(name) for name in facts.interface.get("stub_procedures") or ()}
             ),
             "aliases": sorted(
                 {remote.alias for remote in assembler.remotes.values()}

From 5d876f24870f801a52e7c126f8b29ded4ee1e708 Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 11:43:53 -0600
Subject: [PATCH 16/73] A sibling's INOUT actual is read; a specific fits with
 its optionals left off

The companion table said which positions a call writes and left the
reading to 'everything else', so an INOUT actual -- lhs, rhs, err_info into
CLUBB's band_solve -- was written and not read on the source side. Each
entry now says what the caller reads: IN, INOUT, UNKNOWN and a buffer OUT.
And a specific of a generic fits a call whose actuals fall between its
required and its total count, since an optional dummy may be left off.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/fortran/interface.py | 17 ++++++++++++++---
 src/recast/fortran/rwset.py     | 19 +++++++++++++++++--
 2 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/src/recast/fortran/interface.py b/src/recast/fortran/interface.py
index be37940..f2a70f3 100644
--- a/src/recast/fortran/interface.py
+++ b/src/recast/fortran/interface.py
@@ -1310,6 +1310,14 @@ def companion_externals(record: dict[str, Any]) -> dict[str, dict[str, Any]]:
             "buffer_positions": [
                 at for at, argument in enumerate(sub["args"]) if argument.get("buffer")
             ],
+            # What the caller reads: IN and INOUT actuals, and a buffer OUT.
+            # An INOUT actual is written *and* read; out_positions alone said
+            # only the first.
+            "read_positions": [
+                at
+                for at, argument in enumerate(sub["args"])
+                if argument["intent"] in ("IN", "INOUT", "UNKNOWN") or argument.get("buffer")
+            ],
         }
     # The sibling's generics too: a call spells the generic (CLUBB's
     # ``zt2zm_api`` over grid_class's specifics), and a name the scope does
@@ -1325,9 +1333,10 @@ def companion_externals(record: dict[str, Any]) -> dict[str, dict[str, Any]]:
         # argument count, and the scope picks by the actuals it sees. A
         # union over specifics of different arity marked the wrong
         # positions -- CLUBB's tridiag_solve, zm2zt_api.
-        arity = {
-            s: len(next(x for x in record["subprograms"] if x["name"] == s)["args"])
-            for s, _ in known
+        signatures = {s: next(x for x in record["subprograms"] if x["name"] == s) for s, _ in known}
+        arity = {s: len(sig["args"]) for s, sig in signatures.items()}
+        required = {
+            s: sum(1 for a in sig["args"] if not a.get("optional")) for s, sig in signatures.items()
         }
         table[generic] = {
             "kind": known[0][1]["kind"],
@@ -1339,8 +1348,10 @@ def companion_externals(record: dict[str, Any]) -> dict[str, dict[str, Any]]:
                 {
                     "name": s,
                     "args": arity[s],
+                    "required": required[s],
                     "out_positions": entry["out_positions"],
                     "buffer_positions": entry.get("buffer_positions", []),
+                    "read_positions": entry.get("read_positions", []),
                 }
                 for s, entry in known
             ],
diff --git a/src/recast/fortran/rwset.py b/src/recast/fortran/rwset.py
index 377ee09..29ab24f 100644
--- a/src/recast/fortran/rwset.py
+++ b/src/recast/fortran/rwset.py
@@ -388,17 +388,32 @@ def call(stmt: Any) -> None:
             if external and external.get("specifics"):
                 # A sibling's generic: the specific this arity reaches, or the
                 # union when none matches exactly.
-                fitting = [x for x in external["specifics"] if x["args"] == len(actuals)]
+                # An optional dummy may be left off: a specific fits when the
+                # actuals fall between its required and its total count.
+                fitting = [
+                    x
+                    for x in external["specifics"]
+                    if x.get("required", x["args"]) <= len(actuals) <= x["args"]
+                ]
                 if fitting:
                     external = fitting[0]
             out_positions = set(external.get("out_positions", [])) if external else set()
             buffers = set(external.get("buffer_positions", [])) if external else set()
+            read_positions = (
+                set(external["read_positions"])
+                if external and "read_positions" in external
+                else None
+            )
             for j, actual in enumerate(actuals):
                 if j in out_positions:
                     write_target(actual)
-                if j not in out_positions or j in buffers:
+                if read_positions is not None:
+                    is_read = j in read_positions
+                else:
                     # A buffer OUT of a sibling is read as well as written:
                     # the emitter passes the caller's storage in (#38).
+                    is_read = j not in out_positions or j in buffers
+                if is_read:
                     reads.update(expr_reads(actual, scope))
             return
 

From 79a61c7be355de9bd953ba11fbdc3c07bc30b2a0 Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 11:56:28 -0600
Subject: [PATCH 17/73] A keyword actual into a sibling's procedure lands on
 its own position

call band_solve( ..., solut, rcond = rcond ) (CLUBB): bound by position the
keyword fell on the dummy before it and was counted a read; the companion
table now carries the dummies' names and the scope binds a keyword actual
by name, as it does for a callee of the module itself.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/fortran/interface.py |  4 +++
 src/recast/fortran/rwset.py     |  7 ++++
 tests/test_fortran_analysis.py  | 60 +++++++++++++++++++++++++++++++++
 3 files changed, 71 insertions(+)

diff --git a/src/recast/fortran/interface.py b/src/recast/fortran/interface.py
index f2a70f3..3ecc988 100644
--- a/src/recast/fortran/interface.py
+++ b/src/recast/fortran/interface.py
@@ -1318,6 +1318,9 @@ def companion_externals(record: dict[str, Any]) -> dict[str, dict[str, Any]]:
                 for at, argument in enumerate(sub["args"])
                 if argument["intent"] in ("IN", "INOUT", "UNKNOWN") or argument.get("buffer")
             ],
+            # So a keyword actual (``rcond = rcond``, CLUBB's band_solve) lands
+            # on its own position and not on whichever comes next.
+            "arg_names": [str(argument["name"]).lower() for argument in sub["args"]],
         }
     # The sibling's generics too: a call spells the generic (CLUBB's
     # ``zt2zm_api`` over grid_class's specifics), and a name the scope does
@@ -1352,6 +1355,7 @@ def companion_externals(record: dict[str, Any]) -> dict[str, dict[str, Any]]:
                     "out_positions": entry["out_positions"],
                     "buffer_positions": entry.get("buffer_positions", []),
                     "read_positions": entry.get("read_positions", []),
+                    "arg_names": entry.get("arg_names", []),
                 }
                 for s, entry in known
             ],
diff --git a/src/recast/fortran/rwset.py b/src/recast/fortran/rwset.py
index 29ab24f..1db027a 100644
--- a/src/recast/fortran/rwset.py
+++ b/src/recast/fortran/rwset.py
@@ -404,7 +404,14 @@ def call(stmt: Any) -> None:
                 if external and "read_positions" in external
                 else None
             )
+            if external and external.get("arg_names"):
+                # Keyword actuals by name, positional ones in order, as for a
+                # callee of this module.
+                formals = [{"name": n} for n in external["arg_names"]]
+                actuals = _bind_actuals({"args": formals}, items)
             for j, actual in enumerate(actuals):
+                if actual is None:
+                    continue
                 if j in out_positions:
                     write_target(actual)
                 if read_positions is not None:
diff --git a/tests/test_fortran_analysis.py b/tests/test_fortran_analysis.py
index 8907d88..0555202 100644
--- a/tests/test_fortran_analysis.py
+++ b/tests/test_fortran_analysis.py
@@ -1894,3 +1894,63 @@ def test_a_call_into_a_stubbed_module_reads_and_writes_nothing(tmp_path: Path) -
     blocks = facts.effects["fortran:budget_mod/tend"]["blocks"]
     reads = {name for block in blocks for name in block.get("reads", [])}
     assert "stats_tmp" not in reads and "x_budget" not in reads
+
+
+KEYWORD_CALLER = """\
+module caller_mod
+  use solve_mod, only: solve
+  implicit none
+  private
+  public :: run
+contains
+  subroutine run( n, m, x, rc )
+    integer, intent(in) :: n, m
+    real, intent(inout) :: x(n, m)
+    real, intent(out) :: rc
+    call solve( n, m, x, rcond = rc )
+  end subroutine run
+end module caller_mod
+"""
+
+SOLVER_WITH_OPTIONAL = """\
+module solve_mod
+  implicit none
+  private
+  public :: solve
+  interface solve
+    module procedure solve_one, solve_many
+  end interface
+contains
+  subroutine solve_one( n, x, rcond )
+    integer, intent(in) :: n
+    real, intent(inout) :: x(n)
+    real, intent(out), optional :: rcond
+    x = 2.0 * x
+    if ( present( rcond ) ) rcond = 1.0
+  end subroutine solve_one
+  subroutine solve_many( n, m, x, rcond )
+    integer, intent(in) :: n, m
+    real, intent(inout) :: x(n, m)
+    real, intent(out), optional :: rcond
+    x = 2.0 * x
+    if ( present( rcond ) ) rcond = 1.0
+  end subroutine solve_many
+end module solve_mod
+"""
+
+
+def test_a_keyword_actual_into_a_siblings_generic_lands_on_its_own_position(tmp_path: Path) -> None:
+    """``call band_solve( ..., solut, rcond = rcond )`` (CLUBB): the keyword
+    names an optional OUT dummy at the end. Bound by position it fell on the
+    dummy before it and was read; bound by name it is written, not read --
+    and the specific is picked with the optional counted."""
+    from recast.fortran.frontend import FortranFrontend
+
+    _write(tmp_path, "solve.f90", SOLVER_WITH_OPTIONAL)
+    _write(tmp_path, "caller.f90", KEYWORD_CALLER)
+    frontend = FortranFrontend()
+    unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:caller_mod")
+    facts = frontend.analyze(unit, tmp_path)
+    (block,) = facts.effects["fortran:caller_mod/run"]["blocks"]
+    assert "rc" in block["writes"] and "rc" not in block["reads"]
+    assert "x" in block["writes"] and "x" in block["reads"]  # INOUT: both

From 27ebf1a66f0a2313dc5fe656aef6af97b40884fc Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 12:34:37 -0600
Subject: [PATCH 18/73] Tier-2 replay: logicals and zero-extent records, plain
 OUT dummies, rank-picked specifics, stubbed calls raise in place

Found by replaying CLUBB's advance_* solvers on bomex recordings:

* dump-replay's parser took only numbers: a logical ``T``/``F`` header
  (l_implemented) was no value, and a zero-extent component (scalar
  tracers under sclr_dim = 0) vanished with its empty body. Both are
  values now.
* The recorder wrote what the objects received; the plain OUT/INOUT
  dummies (wp2, wp3, mono_flux_limiter's low_lev_effect) are outputs too.
* Two specifics of one arity (tridiag_solve with its optional rcond and
  the multiple-rhs one without) were picked by count; the ranks of the
  bare-name actuals decide, so an IN logical is no longer scored written.
* A component name on an OUT actual (``pdf_params%chi_1``) was counted a
  read of a variable of that name; the tidier answer the pipeline's test
  kept as a documented disagreement is now the answer on both sides.
* A call to a procedure of a stubbed module that no statement stub
  answers (lapack_band_solvex on CLUBB's LAPACK path) deferred its whole
  block, condition and LU branch included, so the candidate raised on the
  path the run takes. It raises on its own line now.
* ``max(2, edsclr_dim)`` sizing a local is a bound Python can spell; a
  comma in a bound is legal inside max/min alone.
* The flat-plan module lookup tolerated a missing record (mypy).

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/fortran/flatten.py             |   2 +-
 src/recast/fortran/interface.py           |   4 +
 src/recast/fortran/rwset.py               |  23 +++-
 src/recast/oracle/dump_replay.py          |  19 +++-
 src/recast/oracle/record.py               |   8 ++
 src/recast/transform/numpy/expressions.py |  41 ++++++--
 src/recast/transform/numpy/statements.py  |  12 +++
 src/recast/transform/numpy/subprograms.py |   4 +
 src/recast/transform/numpy/translate.py   |   1 +
 tests/test_dump_replay.py                 |  20 ++++
 tests/test_flatten.py                     |  18 ++++
 tests/test_fortran_analysis.py            | 121 ++++++++++++++++++++--
 tests/test_loud_refusals.py               |   8 +-
 tests/test_numpy_translate.py             |  75 ++++++++++++++
 14 files changed, 329 insertions(+), 27 deletions(-)

diff --git a/src/recast/fortran/flatten.py b/src/recast/fortran/flatten.py
index 9adfb48..3ce51b0 100644
--- a/src/recast/fortran/flatten.py
+++ b/src/recast/fortran/flatten.py
@@ -362,7 +362,7 @@ def _state_declaration(name: str, root: Path) -> tuple[str, str] | None:
                 continue
             derived = DERIVED.match(str(entry.get("dtype", "")))
             if derived:
-                return derived.group(1).lower(), str(record.get("module", "")).lower()
+                return derived.group(1).lower(), str((record or {}).get("module", "")).lower()
     return None
 
 
diff --git a/src/recast/fortran/interface.py b/src/recast/fortran/interface.py
index 3ecc988..6f6d305 100644
--- a/src/recast/fortran/interface.py
+++ b/src/recast/fortran/interface.py
@@ -1356,6 +1356,10 @@ def companion_externals(record: dict[str, Any]) -> dict[str, dict[str, Any]]:
                     "buffer_positions": entry.get("buffer_positions", []),
                     "read_positions": entry.get("read_positions", []),
                     "arg_names": entry.get("arg_names", []),
+                    # Two specifics of one arity (CLUBB's tridiag_solve
+                    # with and without its optional rcond) are told apart
+                    # by the ranks of their dummies.
+                    "ranks": [len(a.get("dims") or []) for a in signatures[s]["args"]],
                 }
                 for s, entry in known
             ],
diff --git a/src/recast/fortran/rwset.py b/src/recast/fortran/rwset.py
index 1db027a..a9afc30 100644
--- a/src/recast/fortran/rwset.py
+++ b/src/recast/fortran/rwset.py
@@ -342,7 +342,15 @@ def _write_actual(actual: Any) -> None:
         """
         if isinstance(actual, f03.Name):
             writes.add(str(actual).lower())
-        elif isinstance(actual, (f03.Part_Ref, f03.Data_Ref)):
+        elif isinstance(actual, f03.Data_Ref):
+            # ``pdf_params%chi_1`` as an OUT actual (CLUBB's pdf_closure):
+            # the object is written; a component's *subscripts* are read,
+            # the component's name is not a variable of this scope.
+            writes.add(str(actual.children[0]).lower())
+            for comp in actual.children[1:]:
+                if isinstance(comp, f03.Part_Ref) and comp.children[1] is not None:
+                    reads.update(expr_reads(comp.children[1], scope))
+        elif isinstance(actual, f03.Part_Ref):
             writes.add(str(actual.children[0]).lower())
             for child in actual.children[1:]:
                 reads.update(expr_reads(child, scope))
@@ -395,6 +403,19 @@ def call(stmt: Any) -> None:
                     for x in external["specifics"]
                     if x.get("required", x["args"]) <= len(actuals) <= x["args"]
                 ]
+                if len(fitting) > 1:
+                    # Several fit the count (one's optional tail is the
+                    # other's required one): the ranks of the bare-name
+                    # actuals decide, where a declared rank is known.
+                    def _agrees(x: dict[str, Any]) -> bool:
+                        return all(
+                            not isinstance(a, f03.Name)
+                            or scope.ranks.get(str(a).lower()) is None
+                            or scope.ranks[str(a).lower()] == r
+                            for a, r in zip(actuals, x.get("ranks", []), strict=False)
+                        )
+
+                    fitting = [x for x in fitting if _agrees(x)] or fitting
                 if fitting:
                     external = fitting[0]
             out_positions = set(external.get("out_positions", [])) if external else set()
diff --git a/src/recast/oracle/dump_replay.py b/src/recast/oracle/dump_replay.py
index 833c4c8..8ad33da 100644
--- a/src/recast/oracle/dump_replay.py
+++ b/src/recast/oracle/dump_replay.py
@@ -122,7 +122,18 @@ def parse_dump(text: str) -> tuple[dict[str, Any], dict[str, Any]]:
     all_integers = True  # every value of the section so far is an integer literal
 
     def flush() -> None:
-        if not name or not values:
+        if not name:
+            return
+        if not values:
+            # A zero-extent record: a component the run never allocated
+            # (CLUBB's scalar-tracer arrays under sclr_dim = 0), written as
+            # ``name(1,88,0)``. An empty array of that shape is a value.
+            if dims_text and any(t.strip() == "0" for t in dims_text.split(",")):
+                empty: list[int] = []
+                for t in dims_text.split(","):
+                    sized = str(metadata.get(t.strip().lower(), t.strip()))
+                    empty.append(int(sized) if sized.lstrip("-").isdigit() else 0)
+                (inputs if target == "INPUT" else outputs)[name] = np.zeros(tuple(empty), order="F")
             return
         # The recorder writes integer arrays with ``i0`` and reals with an
         # exponent, so a section whose every value is an integer literal is an
@@ -179,6 +190,12 @@ def flush() -> None:
                         metadata[key] = whole
                         inputs[key] = np.int32(whole)
                 except ValueError:
+                    # A logical, written ``T`` / ``F`` by the recorder's l1
+                    # format: a value, not a diagnostic.
+                    if text_value.upper() in ("T", "F", ".TRUE.", ".FALSE."):
+                        metadata[key] = text_value.upper().startswith(("T", ".T"))
+                        inputs[key] = np.bool_(metadata[key])
+                        continue
                     # Not a number. Upstream swallows this with a bare
                     # ``except``; narrowed to what can actually be raised
                     # here, which changes no outcome and stops the clause
diff --git a/src/recast/oracle/record.py b/src/recast/oracle/record.py
index 535e1de..35f3b0c 100644
--- a/src/recast/oracle/record.py
+++ b/src/recast/oracle/record.py
@@ -302,6 +302,14 @@ def recorder_module(
                 _dims_text(state.extents),
             )
         lines += ["    else", f"       if (n_{sname} > max_calls) return"]
+        # The plain OUT and INOUT dummies are outputs too (CLUBB's advance_*
+        # return wp2, wp3, ... beside what they write into the objects).
+        for a in originals:
+            if DERIVED.match(str(a["dtype"])) or a["intent"] not in ("OUT", "INOUT"):
+                continue
+            _record_call(
+                lines, u, "OUTPUT", a["name"], a["name"], len(a.get("dims") or ()), str(a["dtype"])
+            )
         for obj in plan.objects:
             for comp in obj.components:
                 if comp.written:
diff --git a/src/recast/transform/numpy/expressions.py b/src/recast/transform/numpy/expressions.py
index c88c9f8..9a7aa54 100644
--- a/src/recast/transform/numpy/expressions.py
+++ b/src/recast/transform/numpy/expressions.py
@@ -59,7 +59,7 @@
 
 DIM_KEYWORD = re.compile(r"dim\s*=\s*", re.I)
 
-BOUND_TOKENS = re.compile(r"[A-Za-z_]\w*\s*%\s*[A-Za-z_]\w*|[A-Za-z_]\w*|\d+|[()+\-*/ ]")
+BOUND_TOKENS = re.compile(r"[A-Za-z_]\w*\s*%\s*[A-Za-z_]\w*|[A-Za-z_]\w*|\d+|[()+\-*/, ]")
 """What a declared bound is allowed to be made of. Bound texts are simple by
 construction; anything richer refuses the statement that needed the bound."""
 
@@ -528,25 +528,46 @@ def extent(match: re.Match[str]) -> str:
         if substituted != text:
             text = substituted
         rendered, position = [], 0
+        opens_intrinsic = False  # the next "(" opens a max/min call
+        calls: list[bool] = []  # per open parenthesis: a max/min call?
         for match in BOUND_TOKENS.finditer(text):
             if match.start() != position:
                 raise NoRule(f"dim expr {text!r}")
             position = match.end()
-            token = match.group(0)
-            if "%" in token:
+            piece = match.group(0)
+            if "%" in piece:
                 # ``bounds%begp`` sizing a local: the component of a dummy,
                 # which is an attribute of the same name on this side.
-                root, component = (t.strip() for t in token.split("%", 1))
+                root, component = (t.strip() for t in piece.split("%", 1))
                 rendered.append(f"{self.names.symbol(root)}.{pysafe(component.lower())}")
-            elif re.match(r"[A-Za-z_]", token):
-                rendered.append(self.names.symbol(token))
-            elif token.isdigit() and token not in ("0", "1", "2"):
-                hoisted = self.names.literals.get(token)
+            elif piece.lower() in ("max", "min") and text[match.end() :].lstrip().startswith("("):
+                # ``max(2, edsclr_dim)`` sizing a local (CLUBB's windm
+                # solver): Python spells the two intrinsics the same way,
+                # and a bound's operands are integers. The only calls a
+                # bound may carry; a comma is legal inside one of them alone.
+                rendered.append(piece.lower())
+                opens_intrinsic = True
+            elif re.match(r"[A-Za-z_]", piece):
+                rendered.append(self.names.symbol(piece))
+            elif piece.isdigit() and piece not in ("0", "1", "2"):
+                hoisted = self.names.literals.get(piece)
                 if hoisted is None:
-                    raise NoRule(f"declared dim literal {token}")
+                    raise NoRule(f"declared dim literal {piece}")
                 rendered.append(hoisted)
+            elif piece == "(":
+                calls.append(opens_intrinsic)
+                opens_intrinsic = False
+                rendered.append(piece)
+            elif piece == ")":
+                if calls:
+                    calls.pop()
+                rendered.append(piece)
+            elif piece == ",":
+                if not (calls and calls[-1]):
+                    raise NoRule(f"dim expr {text!r}")
+                rendered.append(piece)
             else:
-                rendered.append(token)
+                rendered.append(piece)
         if position != len(text):
             raise NoRule(f"dim expr {text!r}")
         return "".join(rendered)
diff --git a/src/recast/transform/numpy/statements.py b/src/recast/transform/numpy/statements.py
index 3065c91..ef3e11c 100644
--- a/src/recast/transform/numpy/statements.py
+++ b/src/recast/transform/numpy/statements.py
@@ -193,6 +193,9 @@ class Statements:
     externals: dict[str, dict[str, Any]] = field(default_factory=dict)
     """Procedures with an audited shim in the externals module."""
 
+    stub_procedures: frozenset[str] = frozenset()
+    """Use-imported from a stubbed module (the frontend's list)."""
+
     call_transforms: dict[str, Any] = field(default_factory=dict)
     """Callee -> a domain package's answer for it; see ``calls.CallSite``."""
 
@@ -1728,6 +1731,15 @@ def _call(self, node: Any, indent: int) -> list[str]:
                 # list says "the engine does not know this intrinsic" rather
                 # than "this tree is missing a library".
                 raise NoRule(f"intrinsic subroutine {name!r} has no rule")
+            if name in self.stub_procedures:
+                # A procedure of a stubbed module that no statement stub
+                # answers (CLUBB's lapack_band_solvex: the LAPACK path the
+                # run does not take). A raise on the statement keeps the
+                # branch around it -- deferring the block dropped the
+                # ``if ( method == lapack )`` with it, and the candidate
+                # raised on the path the run *does* take.
+                reason = f"{name}: procedure of a stubbed module, not ported"
+                return [f"{pad}raise NotImplementedError({reason!r})"]
             raise NoRule(f"call to external subroutine {name!r}")
 
         # Bind actuals to formals BY NAME for keyword arguments: Fortran
diff --git a/src/recast/transform/numpy/subprograms.py b/src/recast/transform/numpy/subprograms.py
index 130918b..992ce3d 100644
--- a/src/recast/transform/numpy/subprograms.py
+++ b/src/recast/transform/numpy/subprograms.py
@@ -149,6 +149,9 @@ class Subprograms:
     companion_globals: dict[str, str] = field(default_factory=dict)
     externals: dict[str, dict[str, Any]] = field(default_factory=dict)
     remotes: dict[str, Remote] = field(default_factory=dict)
+    stub_procedures: frozenset[str] = frozenset()
+    """Procedures use-imported from a stubbed module; a call to one that no
+    statement stub answers is a raise, not a deferral of its block."""
     function_stubs: dict[str, str] = field(default_factory=dict)
     statement_stubs: dict[str, str] = field(default_factory=dict)
     intrinsics: dict[str, Any] = field(default_factory=dict)
@@ -333,6 +336,7 @@ def floors(self, name: str) -> Statements:
             names,
             expressions,
             externals=self.externals,
+            stub_procedures=self.stub_procedures,
             stubs=dict(self.statement_stubs),
             call_transforms=dict(self.call_transforms),
             poison_undefined=self.poison_undefined,
diff --git a/src/recast/transform/numpy/translate.py b/src/recast/transform/numpy/translate.py
index 368e6d6..1842897 100644
--- a/src/recast/transform/numpy/translate.py
+++ b/src/recast/transform/numpy/translate.py
@@ -280,6 +280,7 @@ def apply(self, unit: Unit, facts: Facts, config: dict[str, Any]) -> Candidate:
             use_parameters=use_parameters,
             companion_globals=companion_globals,
             externals=facts.provenance.get("externals", {}),
+            stub_procedures=frozenset(facts.interface.get("stub_procedures") or ()),
             remotes=remotes,
             function_stubs=config.get("function_stubs", {}),
             statement_stubs=config.get("statement_stubs", {}),
diff --git a/tests/test_dump_replay.py b/tests/test_dump_replay.py
index 75df905..b5a6c87 100644
--- a/tests/test_dump_replay.py
+++ b/tests/test_dump_replay.py
@@ -361,3 +361,23 @@ def test_the_shipped_example_replays_bit_exact(tmp_path: Path) -> None:
     assert set(compared) == {"settle", "column_mass"}
     assert all(outcome["points"] > 0 for outcome in compared.values())
     assert verdict.metrics["uncovered"] == []
+
+
+def test_a_logical_header_scalar_is_an_input() -> None:
+    """CLUBB's ``l_implemented``: the recorder writes a logical ``T``/``F``.
+    The parser took only numbers, so the replay had no value for it."""
+    inputs, _ = parse_dump("# PROBE m.s: call=1\n# l_on = T\n# l_off = F\n# INPUT: x(1)\n1.0\n")
+    assert inputs["l_on"] is not None and bool(inputs["l_on"]) is True
+    assert bool(inputs["l_off"]) is False
+    assert inputs["l_on"].dtype == np.bool_
+
+
+def test_a_zero_extent_array_is_a_value() -> None:
+    """A component the run never allocated (CLUBB's scalar tracers under
+    ``sclr_dim = 0``) is written ``name(1,3,0)`` with nothing under it.
+    Dropped, the replay said the record carried no value for it."""
+    text = "# PROBE m.s: call=1\n# INPUT: s(1,3,0)\n# OUTPUT: t(0)\n# OUTPUT: y(1)\n2.0\n"
+    inputs, outputs = parse_dump(text)
+    assert inputs["s"].shape == (1, 3, 0)
+    assert outputs["t"].shape == (0,)
+    assert outputs["y"].tolist() == [2.0]
diff --git a/tests/test_flatten.py b/tests/test_flatten.py
index c557a27..f0c451d 100644
--- a/tests/test_flatten.py
+++ b/tests/test_flatten.py
@@ -688,3 +688,21 @@ def test_the_recorder_guards_a_component_the_run_may_not_allocate(tmp_path: Path
     assert "if (allocated(c%coef)) then" in recorder
     assert "'# INPUT: c__coef(0,0)'" in recorder
     assert "merge(size(c%coef, 2), 0, allocated(c%coef))" in recorder
+
+
+def test_the_recorder_writes_the_plain_out_dummies_too(tmp_path: Path) -> None:
+    """CLUBB's advance_* hand back ``wp2``, ``wp3``... as INOUT dummies beside
+    what they write into their objects. Recorded only through the objects,
+    the replay found no value for the required outputs (mono_flux_limiter's
+    ``low_lev_effect``)."""
+    from recast.oracle.record import recorder_module
+
+    (tmp_path / "coefs_mod.f90").write_text(COEFS)
+    (tmp_path / "solver_mod.f90").write_text(USES_COEFS)
+    frontend = FortranFrontend(flatten={"patch_count": "ngrdcol", "bounds_pattern": r"^ngrdcol$"})
+    unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:solver_mod")
+    facts = frontend.analyze(unit, tmp_path)
+    (plan,) = plans_for(facts, tmp_path, CLUBB_CONVENTIONS)
+    recorder = recorder_module("solver_mod", [plan])
+    assert "call rec_r1(u_apply, 'INPUT', 'x', trim(dims), reshape(x, (/size(x)/)))" in recorder
+    assert "call rec_r1(u_apply, 'OUTPUT', 'x', trim(dims), reshape(x, (/size(x)/)))" in recorder
diff --git a/tests/test_fortran_analysis.py b/tests/test_fortran_analysis.py
index 0555202..db17f35 100644
--- a/tests/test_fortran_analysis.py
+++ b/tests/test_fortran_analysis.py
@@ -726,15 +726,16 @@ def test_a_call_splits_its_arguments_by_declared_intent(tmp_path: Path) -> None:
 """
 
 
-def test_a_component_name_is_read_on_the_out_argument_path_only(tmp_path: Path) -> None:
+def test_a_component_name_is_not_a_read_on_the_out_argument_path(tmp_path: Path) -> None:
     """``b % q`` writes ``b``. On an assignment, ``q`` is an attribute and not
     a symbol; passed to an intent(out) dummy, the pipeline this came from
-    counts it as a read as well.
+    counted it as a read as well.
 
-    The two disagree, and the disagreement is preserved. Resolving it would be
-    a change to answers a bit-exact gate has been run against, and the two
-    sites in CAM where it shows are both in modules with no translation to
-    check the tidier answer against.
+    The disagreement was preserved until a translation showed the tidier
+    answer: CLUBB's pdf_closure passes ``pdf_params%chi_1`` and six more
+    components as OUT actuals, the candidate spells attributes and reads no
+    variable of those names, and the gate scored six blocks as disagreeing
+    over reads of variables the scope does not have.
     """
     from recast.fortran import rwset
 
@@ -747,10 +748,8 @@ def test_a_component_name_is_read_on_the_out_argument_path_only(tmp_path: Path)
     )
     blocks = {b["id"]: b for b in rwset.block_rwsets(node, rwset.scope_for(record, "drive"))}
     # ``slot(:)`` is a caller-buffer OUT (#36), so ``b`` is read as well as
-    # written on the call (#38); ``q`` is the pipeline's component read.
-    assert blocks["B001"] == {"id": "B001", "reads": ["b", "n", "q"], "writes": ["b"]}, (
-        "out-argument"
-    )
+    # written on the call (#38); ``q`` is an attribute on both paths.
+    assert blocks["B001"] == {"id": "B001", "reads": ["b", "n"], "writes": ["b"]}, "out-argument"
     assert blocks["B002"] == {"id": "B002", "reads": ["n"], "writes": ["b"]}, "assignment"
 
 
@@ -1954,3 +1953,105 @@ def test_a_keyword_actual_into_a_siblings_generic_lands_on_its_own_position(tmp_
     (block,) = facts.effects["fortran:caller_mod/run"]["blocks"]
     assert "rc" in block["writes"] and "rc" not in block["reads"]
     assert "x" in block["writes"] and "x" in block["reads"]  # INOUT: both
+
+
+RANK_OVERLOADED_SOLVER = """
+module solve_mod
+  implicit none
+  private
+  public :: solve
+  interface solve
+    module procedure solve_one, solve_many
+  end interface solve
+contains
+  subroutine solve_one( n, flag, a, x, rc )
+    integer, intent(in) :: n
+    logical, intent(in) :: flag
+    real, intent(inout) :: a(n)
+    real, intent(out) :: x(n)
+    real, intent(out), optional :: rc
+    x = a
+    if ( present(rc) ) rc = 1.0
+  end subroutine solve_one
+  subroutine solve_many( n, m, flag, a, x )
+    integer, intent(in) :: n, m
+    logical, intent(in) :: flag
+    real, intent(inout) :: a(n, m)
+    real, intent(out) :: x(n, m)
+    x = a
+  end subroutine solve_many
+end module solve_mod
+"""
+
+RANK_CALLER = """
+module caller_mod
+  use solve_mod, only: solve
+  implicit none
+contains
+  subroutine run( n, m, flag, a, x )
+    integer, intent(in) :: n, m
+    logical, intent(in) :: flag
+    real, intent(inout) :: a(n, m)
+    real, intent(out) :: x(n, m)
+    call solve( n, m, flag, a, x )
+  end subroutine run
+end module caller_mod
+"""
+
+
+def test_specifics_of_one_arity_are_told_apart_by_rank(tmp_path: Path) -> None:
+    """CLUBB's ``tridiag_solve``: the single-rhs specific with its optional
+    ``rcond`` takes as many actuals as the multiple-rhs one without. Picked
+    by count alone the first won, and ``l_implemented`` -- an IN logical on
+    the position where the other specific's ``rhs`` sits -- was scored
+    written. The ranks of the actuals pick the specific whose dummies match."""
+    from recast.fortran.frontend import FortranFrontend
+
+    _write(tmp_path, "solve.f90", RANK_OVERLOADED_SOLVER)
+    _write(tmp_path, "caller.f90", RANK_CALLER)
+    frontend = FortranFrontend()
+    unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:caller_mod")
+    facts = frontend.analyze(unit, tmp_path)
+    (block,) = facts.effects["fortran:caller_mod/run"]["blocks"]
+    assert "flag" not in block["writes"] and "flag" in block["reads"]
+    assert "x" in block["writes"] and "x" not in block["reads"]
+    assert "a" in block["writes"] and "a" in block["reads"]
+
+
+COMPONENT_OUT_CALLER = """
+module pdf_mod
+  implicit none
+  type pdf_type
+    real, allocatable :: chi(:), eta(:)
+  end type pdf_type
+contains
+  subroutine fill( n, chi, eta )
+    integer, intent(in) :: n
+    real, intent(in) :: chi(n)
+    real, intent(out) :: eta(n)
+    eta = chi
+  end subroutine fill
+  subroutine run( n, p )
+    integer, intent(in) :: n
+    type(pdf_type), intent(inout) :: p
+    call fill( n, p%chi, p%eta )
+  end subroutine run
+end module pdf_mod
+"""
+
+
+def test_a_component_out_actual_reads_no_variable_of_the_components_name(tmp_path: Path) -> None:
+    """``call fill( n, p%chi, p%eta )`` (CLUBB's pdf_closure passes
+    ``pdf_params%chi_1`` and friends as OUT actuals): the object is written.
+    The component's bare name was counted as a read of a variable ``eta``
+    that the scope does not have; the translation, spelling the attribute,
+    read no such thing, and the block disagreed."""
+    from recast.fortran.frontend import FortranFrontend
+
+    _write(tmp_path, "pdf.f90", COMPONENT_OUT_CALLER)
+    frontend = FortranFrontend()
+    unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:pdf_mod")
+    facts = frontend.analyze(unit, tmp_path)
+    (block,) = facts.effects["fortran:pdf_mod/run"]["blocks"]
+    assert "p" in block["writes"] and "p" in block["reads"]
+    assert "eta" not in block["reads"] and "chi" not in block["reads"]
diff --git a/tests/test_loud_refusals.py b/tests/test_loud_refusals.py
index 6cc3621..83fc9bf 100644
--- a/tests/test_loud_refusals.py
+++ b/tests/test_loud_refusals.py
@@ -44,10 +44,10 @@
 
   subroutine prologue_refusals(n, out1, w)
     integer, intent(in) :: n
-    real(r8), intent(out) :: out1(max(n, 2))
+    real(r8), intent(out) :: out1(mod(n, 3) + 2)
     type(wide_t), intent(out) :: w
     integer, parameter :: grid(2, 2) = reshape((/ 1, 2, 3, 4 /), (/ 2, 2 /))
-    real(r8) :: scr(max(n, 2))
+    real(r8) :: scr(mod(n, 3) + 2)
     out1 = 0.0_r8
     scr = 0.0_r8
     w%rows = 0.0_r8
@@ -137,10 +137,10 @@ def test_every_prologue_refusal_is_recorded_and_raises(source: Path, renderer: M
     prologue = [entry for entry in deferred if entry["block"].startswith("P")]
     assert [entry["block"] for entry in prologue] == ["P001", "P002", "P003", "P004"]
     reasons = "\n".join(entry["reason"] for entry in prologue)
-    assert "out-arg out1: allocation refused (dim expr 'MAX(n, 2)')" in reasons
+    assert "out-arg out1: allocation refused (dim expr 'MOD(n, 3) + 2')" in reasons
     assert "out-arg w: INTENT(OUT) derived-type dummy not materialized" in reasons
     assert "local parameter grid" in reasons
-    assert "local array scr: extent not resolvable (dim expr 'MAX(n, 2)')" in reasons
+    assert "local array scr: extent not resolvable (dim expr 'MOD(n, 3) + 2')" in reasons
     # The old wording is gone, and every refusal raises.
     assert "allocation skipped" not in body
     assert "prologue skipped" not in body
diff --git a/tests/test_numpy_translate.py b/tests/test_numpy_translate.py
index 22054bc..73e8f44 100644
--- a/tests/test_numpy_translate.py
+++ b/tests/test_numpy_translate.py
@@ -355,3 +355,78 @@ def test_a_loop_index_read_after_the_loop_has_the_completion_value(tmp_path: Pat
     finally:
         sys.path.remove(str(out))
         sys.modules.pop("search_mod_numpy", None)
+
+
+STUBBED_PATH = """
+module lapack_wrap
+  implicit none
+contains
+  subroutine band_solvex( n, a, x )
+    integer, intent(in) :: n
+    real, intent(inout) :: a(n)
+    real, intent(out) :: x(n)
+    x = a
+  end subroutine band_solvex
+end module lapack_wrap
+"""
+
+CHOOSES_A_SOLVER = """
+module solver_mod
+  use lapack_wrap, only: band_solvex
+  implicit none
+contains
+  subroutine solve( method, n, m, a, x )
+    integer, intent(in) :: method, n, m
+    real, intent(inout) :: a(n)
+    real, intent(out) :: x(n)
+    real :: work(n, max(2, m))
+    work = 0.0
+    if ( method == 1 ) then
+      call band_solvex( n, a, x )
+    else
+      x = a + work(:, 1)
+    end if
+  end subroutine solve
+end module solver_mod
+"""
+
+
+def test_a_call_into_a_stubbed_module_raises_on_its_own_line(tmp_path: Path) -> None:
+    """CLUBB's matrix_solver_wrapper chooses LAPACK or its own LU solver by
+    a run-time flag; ``lapack_wrap`` is stubbed. With no rule for the call
+    the whole IF was deferred -- condition and LU branch included -- and the
+    candidate raised on the path the run takes. The raise belongs to the
+    statement; the branch around it stays. And ``work(n, max(2, m))``
+    (windm's ``rhs``) is a bound Python can spell."""
+    import importlib
+    import sys
+
+    (tmp_path / "lapack_wrap.f90").write_text(STUBBED_PATH)
+    (tmp_path / "solver_mod.f90").write_text(CHOOSES_A_SOLVER)
+    frontend = FortranFrontend(stub_modules=["lapack_wrap"])
+    unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:solver_mod")
+    facts = frontend.analyze(unit, tmp_path)
+    assert facts.interface["stub_procedures"] == ["band_solvex"]
+    candidate = NumpyTranslation().apply(unit, facts, {"root": tmp_path})
+    assert not candidate.deferred
+    out = tmp_path / "emitted"
+    out.mkdir()
+    for path, content in candidate.files.items():
+        (out / path.name).write_bytes(content)
+    text = (out / "solver_mod_numpy.py").read_text()
+    assert "max(2, m)" in text
+    raised = "raise NotImplementedError('band_solvex: procedure of a stubbed module, not ported')"
+    assert raised in text
+    sys.path.insert(0, str(out))
+    try:
+        module = importlib.import_module("solver_mod_numpy")
+        import numpy as np
+
+        a = np.array([1.0, 2.0], dtype=np.float32)
+        _a, x = module.solve(2, 2, 1, a)  # INOUT a and OUT x come back
+        assert x.tolist() == [1.0, 2.0]
+        with pytest.raises(NotImplementedError):
+            module.solve(1, 2, 1, a)
+    finally:
+        sys.path.remove(str(out))
+        sys.modules.pop("solver_mod_numpy", None)

From 4d3ac1a1391af5e977c1d45540c1a698c04a4899 Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 12:37:04 -0600
Subject: [PATCH 19/73] An optional OUT handed on carries the caller's presence

``call inner( x, y, rc = rc )`` where ``rc`` is the caller's own optional
OUT was rendered ``want_rc=True``: present in the callee on every call.
CLUBB's xm_wpxp_solve hands ``rcond = rcond`` to band_solve, whose
``present(rcond)`` branch is the LAPACK diagnostic path -- taken, on the
replay, on every call the run made without rcond. The presence is the
caller's own sentinel for an optional OUT and ``is not None`` for an
optional IN.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/numpy/statements.py | 17 ++++++-
 tests/test_numpy_translate.py            | 57 ++++++++++++++++++++++++
 2 files changed, 73 insertions(+), 1 deletion(-)

diff --git a/src/recast/transform/numpy/statements.py b/src/recast/transform/numpy/statements.py
index ef3e11c..344e71c 100644
--- a/src/recast/transform/numpy/statements.py
+++ b/src/recast/transform/numpy/statements.py
@@ -1782,7 +1782,7 @@ def _call(self, node: Any, indent: int) -> list[str]:
                     outputs.append("_")  # the return tuple has fixed length
                 continue
             if self.is_optional_output(formal):
-                inputs.append(f"want_{formal['name']}=True")
+                inputs.append(f"want_{formal['name']}={self._presence(actual)}")
             passes = formal["intent"] in ("IN", "INOUT", "UNKNOWN") or bool(
                 formal.get("buffer") and self.buffer_out_arrays
             )
@@ -1932,6 +1932,21 @@ def _external_call(self, name: str, external: dict[str, Any], node: Any, pad: st
             return [f"{pad}{', '.join(outputs)} = {call}"]
         return [f"{pad}{call}"]
 
+    def _presence(self, actual: Any) -> str:
+        """``present()`` of what is passed to an optional OUT: true for a
+        value of the caller's own; the caller's own presence when the
+        actual is one of its optional dummies handed on (CLUBB's
+        xm_wpxp_solve passes ``rcond = rcond`` to band_solve, and takes the
+        LAPACK diagnostic path only when *its* caller asked for rcond)."""
+        if isinstance(actual, f03.Name):
+            name = str(actual).lower()
+            for declared in self.semantics.subprogram["args"]:
+                if declared["name"].lower() == name and declared["optional"]:
+                    if self.is_optional_output(declared):
+                        return f"want_{name}"
+                    return f"({self.names.symbol(name)} is not None)"
+        return "True"
+
     @staticmethod
     def is_optional_output(formal: dict[str, Any]) -> bool:
         """Optional OUT, the ``want_<name>`` sentinel convention: a trailing
diff --git a/tests/test_numpy_translate.py b/tests/test_numpy_translate.py
index 73e8f44..154269f 100644
--- a/tests/test_numpy_translate.py
+++ b/tests/test_numpy_translate.py
@@ -430,3 +430,60 @@ def test_a_call_into_a_stubbed_module_raises_on_its_own_line(tmp_path: Path) ->
     finally:
         sys.path.remove(str(out))
         sys.modules.pop("solver_mod_numpy", None)
+
+
+HANDS_ON_AN_OPTIONAL = """
+module relay_mod
+  implicit none
+contains
+  subroutine inner( x, y, rc )
+    real, intent(in) :: x
+    real, intent(out) :: y
+    real, intent(out), optional :: rc
+    y = 2.0 * x
+    if ( present(rc) ) rc = 1.0 / x
+  end subroutine inner
+  subroutine outer( x, y, rc, scale )
+    real, intent(in) :: x
+    real, intent(out) :: y
+    real, intent(out), optional :: rc
+    real, intent(in), optional :: scale
+    call inner( x, y, rc = rc )
+    if ( present(scale) ) y = y * scale
+  end subroutine outer
+end module relay_mod
+"""
+
+
+def test_an_optional_handed_on_carries_its_own_presence(tmp_path: Path) -> None:
+    """``call inner( x, y, rc = rc )`` where ``rc`` is the caller's own
+    optional OUT: present in the callee exactly when present in the caller.
+    Rendered ``want_rc=True`` it was always present, and CLUBB's
+    xm_wpxp_solve took the LAPACK diagnostic path on every call."""
+    import importlib
+    import sys
+
+    (tmp_path / "relay_mod.f90").write_text(HANDS_ON_AN_OPTIONAL)
+    frontend = FortranFrontend()
+    unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:relay_mod")
+    facts = frontend.analyze(unit, tmp_path)
+    candidate = NumpyTranslation().apply(unit, facts, {"root": tmp_path})
+    out = tmp_path / "emitted"
+    out.mkdir()
+    for path, content in candidate.files.items():
+        (out / path.name).write_bytes(content)
+    text = (out / "relay_mod_numpy.py").read_text()
+    assert "want_rc=want_rc" in text
+    assert "want_rc=True" not in text
+    sys.path.insert(0, str(out))
+    try:
+        module = importlib.import_module("relay_mod_numpy")
+        import numpy as np
+
+        y, rc = module.outer(np.float32(4.0))
+        assert y == 8.0 and rc != 0.25  # not asked for, not computed
+        y, rc = module.outer(np.float32(4.0), want_rc=True)
+        assert y == 8.0 and rc == 0.25
+    finally:
+        sys.path.remove(str(out))
+        sys.modules.pop("relay_mod_numpy", None)

From 866743eb3ae15ea958d7f3ea85f93b827e3c7374 Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 12:46:05 -0600
Subject: [PATCH 20/73] An optional handed on to an optional OUT is read for
 its presence

The source side of the read/write gate scored call solve( ..., rc = rc )
as a write of the caller's own optional rc; the translation spells the
callee's present(rc) as the caller's want_rc sentinel, a read. Both
sides now count the read, for a sibling's procedure (companion externals
carry the optional OUT positions) and for one of this module.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/fortran/interface.py |  8 ++++
 src/recast/fortran/rwset.py     | 18 +++++++++
 tests/test_fortran_analysis.py  | 66 +++++++++++++++++++++++++++++++++
 3 files changed, 92 insertions(+)

diff --git a/src/recast/fortran/interface.py b/src/recast/fortran/interface.py
index 6f6d305..fed5b24 100644
--- a/src/recast/fortran/interface.py
+++ b/src/recast/fortran/interface.py
@@ -1321,6 +1321,13 @@ def companion_externals(record: dict[str, Any]) -> dict[str, dict[str, Any]]:
             # So a keyword actual (``rcond = rcond``, CLUBB's band_solve) lands
             # on its own position and not on whichever comes next.
             "arg_names": [str(argument["name"]).lower() for argument in sub["args"]],
+            # Handing the caller's own optional to one of these queries its
+            # presence -- a read of it, as ``present()`` is.
+            "optional_out_positions": [
+                at
+                for at, argument in enumerate(sub["args"])
+                if argument.get("optional") and argument["intent"] == "OUT"
+            ],
         }
     # The sibling's generics too: a call spells the generic (CLUBB's
     # ``zt2zm_api`` over grid_class's specifics), and a name the scope does
@@ -1356,6 +1363,7 @@ def companion_externals(record: dict[str, Any]) -> dict[str, dict[str, Any]]:
                     "buffer_positions": entry.get("buffer_positions", []),
                     "read_positions": entry.get("read_positions", []),
                     "arg_names": entry.get("arg_names", []),
+                    "optional_out_positions": entry.get("optional_out_positions", []),
                     # Two specifics of one arity (CLUBB's tridiag_solve
                     # with and without its optional rcond) are told apart
                     # by the ranks of their dummies.
diff --git a/src/recast/fortran/rwset.py b/src/recast/fortran/rwset.py
index a9afc30..2b5159d 100644
--- a/src/recast/fortran/rwset.py
+++ b/src/recast/fortran/rwset.py
@@ -388,6 +388,19 @@ def call(stmt: Any) -> None:
         )
         if name in dummies:
             reads.add(name)
+        optional_dummies = (
+            {a["name"].lower() for a in scope.semantics.subprogram["args"] if a.get("optional")}
+            if (scope.semantics is not None)
+            else set()
+        )
+
+        def hands_on_presence(actual: Any) -> None:
+            """The caller's own optional passed to an optional OUT (CLUBB's
+            xm_wpxp_solve hands ``rcond = rcond`` to band_solve): the callee
+            asks whether it is present, which is a read of it on both sides
+            -- ``present(x)`` here, the ``want_x`` sentinel there."""
+            if isinstance(actual, f03.Name) and str(actual).lower() in optional_dummies:
+                reads.add(str(actual).lower())
 
         if callee is None:
             external = scope.externals.get(name)
@@ -430,11 +443,14 @@ def _agrees(x: dict[str, Any]) -> bool:
                 # callee of this module.
                 formals = [{"name": n} for n in external["arg_names"]]
                 actuals = _bind_actuals({"args": formals}, items)
+            optional_out = set(external.get("optional_out_positions", [])) if external else set()
             for j, actual in enumerate(actuals):
                 if actual is None:
                     continue
                 if j in out_positions:
                     write_target(actual)
+                if j in optional_out:
+                    hands_on_presence(actual)
                 if read_positions is not None:
                     is_read = j in read_positions
                 else:
@@ -457,6 +473,8 @@ def _agrees(x: dict[str, Any]) -> bool:
                 reads.update(expr_reads(actual, scope))
             if formal["intent"] in ("OUT", "INOUT"):
                 _write_actual(actual)
+            if formal.get("optional") and formal["intent"] == "OUT":
+                hands_on_presence(actual)
 
     def visit(stmt: Any) -> None:
         if isinstance(stmt, f08.Block_Construct):
diff --git a/tests/test_fortran_analysis.py b/tests/test_fortran_analysis.py
index db17f35..826912e 100644
--- a/tests/test_fortran_analysis.py
+++ b/tests/test_fortran_analysis.py
@@ -2055,3 +2055,69 @@ def test_a_component_out_actual_reads_no_variable_of_the_components_name(tmp_pat
     (block,) = facts.effects["fortran:pdf_mod/run"]["blocks"]
     assert "p" in block["writes"] and "p" in block["reads"]
     assert "eta" not in block["reads"] and "chi" not in block["reads"]
+
+
+HANDS_ON_SOLVER = """
+module solve_mod
+  implicit none
+contains
+  subroutine solve( n, a, x, rc )
+    integer, intent(in) :: n
+    real, intent(in) :: a(n)
+    real, intent(out) :: x(n)
+    real, intent(out), optional :: rc
+    x = a
+    if ( present(rc) ) rc = 1.0
+  end subroutine solve
+end module solve_mod
+"""
+
+HANDS_ON_CALLER = """
+module relay_mod
+  use solve_mod, only: solve
+  implicit none
+contains
+  subroutine outer( n, a, x, rc )
+    integer, intent(in) :: n
+    real, intent(in) :: a(n)
+    real, intent(out) :: x(n)
+    real, intent(out), optional :: rc
+    call solve( n, a, x, rc = rc )
+  end subroutine outer
+  subroutine own( n, a, x, rc )
+    integer, intent(in) :: n
+    real, intent(in) :: a(n)
+    real, intent(out) :: x(n)
+    real, intent(out), optional :: rc
+    call inner( n, a, x, rc )
+  end subroutine own
+  subroutine inner( n, a, x, rc )
+    integer, intent(in) :: n
+    real, intent(in) :: a(n)
+    real, intent(out) :: x(n)
+    real, intent(out), optional :: rc
+    x = a
+    if ( present(rc) ) rc = 2.0
+  end subroutine inner
+end module relay_mod
+"""
+
+
+def test_an_optional_handed_on_to_an_optional_out_is_read_for_its_presence(tmp_path: Path) -> None:
+    """``call solve( ..., rc = rc )`` with the caller's own optional ``rc``:
+    the callee asks ``present(rc)``, which the translation spells as the
+    caller's ``want_rc`` sentinel -- a read of ``rc`` on the target side.
+    The source side scored only the write, and CLUBB's xm_wpxp_solve
+    disagreed on the one block that hands ``rcond`` to band_solve. Both a
+    sibling's procedure and one of this module count it."""
+    from recast.fortran.frontend import FortranFrontend
+
+    _write(tmp_path, "solve.f90", HANDS_ON_SOLVER)
+    _write(tmp_path, "relay.f90", HANDS_ON_CALLER)
+    frontend = FortranFrontend()
+    unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:relay_mod")
+    facts = frontend.analyze(unit, tmp_path)
+    (block,) = facts.effects["fortran:relay_mod/outer"]["blocks"]
+    assert "rc" in block["reads"] and "rc" in block["writes"]
+    (block,) = facts.effects["fortran:relay_mod/own"]["blocks"]
+    assert "rc" in block["reads"] and "rc" in block["writes"]

From b8521607ea8f253d906eb4e78ba0d0dea70dc256 Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 13:06:21 -0600
Subject: [PATCH 21/73] Module state handed whole to a function is carried; a
 constant expression over a constructor is a value (#26)

* The flat plan followed only the caller's own dummies into callees. A
  module-state object passed whole to a derived-type dummy -- CLUBB's
  advance_xm_wpxp hands sponge_layer_damping's profile to sponge_damp_xm,
  which reads ``profile%tau_sponge_damp`` -- was left to the module, and
  the replay found it unallocated. It is followed under the callee's
  dummy now, so its components come back under the state's own name.
  A derived-type state variable is no longer listed as left to the
  module: it is the objects path's business.
* ``100._core_rknd * (/ ... /)`` as a local parameter (#26, saturation and
  pdf_closure's hybrid-PDF path): the token pass rendered a bare
  constructor and handed anything around one to the parser, whose
  literals were never hoisted. Each constructor is rendered where it
  stands and the rest goes through the token pass, elementwise.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/fortran/flatten.py             | 28 +++++++--
 src/recast/transform/numpy/subprograms.py | 31 ++++++++--
 tests/test_flatten.py                     | 73 +++++++++++++++++++++++
 tests/test_numpy_translate.py             | 56 +++++++++++++++++
 4 files changed, 177 insertions(+), 11 deletions(-)

diff --git a/src/recast/fortran/flatten.py b/src/recast/fortran/flatten.py
index 3ce51b0..100ef42 100644
--- a/src/recast/fortran/flatten.py
+++ b/src/recast/fortran/flatten.py
@@ -671,6 +671,20 @@ def take(stmt: Any) -> None:
                     spelled = aliases[spelled].split("%", 1)[0]
                 if spelled in objects or spelled in interesting:
                     mapping[dummy] = spelled
+                elif (
+                    spelled not in declared_here
+                    and spelled not in procedures
+                    and any(
+                        a["name"].lower() == dummy and DERIVED.match(str(a.get("dtype", "")))
+                        for a in callee_record["args"]
+                    )
+                ):
+                    # A module-state object handed whole to a derived-type
+                    # dummy (CLUBB passes sponge_layer_damping's profile to
+                    # its sponge_damp_xm): followed under the callee's
+                    # dummy, so the components the callee reads come back
+                    # under the state's own name and the plan carries them.
+                    mapping[dummy] = spelled
             # A procedure passed as an actual (``hybrid(..., func, ...)``) is
             # called back with the object; its own derived-type dummies are
             # mapped by name, which is how the model spells them.
@@ -1068,17 +1082,21 @@ def _state_vars(
             # the tree or not -- is the run's to say; parameters are not here.
             if name not in wanted or name in seen:
                 continue
+            if (
+                DERIVED.match(str(entry.get("dtype")))
+                or str(entry.get("dtype")) not in FORTRAN_TYPES
+            ):
+                # A derived-type state variable is an *object* of the plan
+                # (carried by component, through the objects path); naming
+                # it here as left to the module misdescribed the sponge
+                # profiles CLUBB's solvers read through a companion.
+                continue
             if default_private and name not in public:
                 # No ``use`` reaches it, so no adapter sets it: both sides
                 # run with the module's own default, and the plan says so.
                 seen.add(name)
                 plan.left_to_module.append(f"{module}%{name}")
                 continue
-            if (
-                DERIVED.match(str(entry.get("dtype")))
-                or str(entry.get("dtype")) not in FORTRAN_TYPES
-            ):
-                continue
             dims = entry.get("dims") or []
             names = [
                 t.lower()
diff --git a/src/recast/transform/numpy/subprograms.py b/src/recast/transform/numpy/subprograms.py
index 992ce3d..b01f79c 100644
--- a/src/recast/transform/numpy/subprograms.py
+++ b/src/recast/transform/numpy/subprograms.py
@@ -84,15 +84,15 @@ def _token_pass_guessed(text: str, spelled: str) -> bool:
     ``_is_expression`` only asks whether Python can *parse* it. Three shapes
     parse and are still wrong: a call or array reference of an uppercased
     name; a ``//`` concatenation, which Python reads as floor division; and an
-    array constructor that was only part of the text (``reshape((/.../),
-    ...)``), where the search silently dropped everything around it.
+    array constructor the rendering dropped.
     """
     if UPPERCASED_CALL.search(spelled):
         return True
     if "//" in text:
         return True
-    constructed = ARRAY_CONSTRUCTOR.search(text)
-    return constructed is not None and constructed.span() != (0, len(text))
+    # A constructor inside a larger text is rendered where it stands now;
+    # the guess is a rendering that lost it.
+    return ARRAY_CONSTRUCTOR.search(text) is not None and "np.array(" not in spelled
 
 
 def _is_expression(text: str) -> str | bool:
@@ -1011,8 +1011,8 @@ def _token_parameter_value(text: str, local_parameters: frozenset[str]) -> str:
             # rather than anything this file would notice.
             exponent = text.replace(" ", "").split("_")[0]
             return "np.float64('" + exponent.replace("d", "e").replace("D", "e") + "')"
-        constructed = ARRAY_CONSTRUCTOR.search(text)
-        if constructed:
+
+        def array_of(constructed: re.Match[str]) -> str:
             items = [strip_kind(item.strip()) for item in constructed.group(1).split(",")]
             if all(re.fullmatch(r"'[^']*'", item) for item in items):
                 return f"np.array([{', '.join(items)}])"
@@ -1030,6 +1030,25 @@ def case_of(match: re.Match[str]) -> str:
             token = match.group()
             return pysafe(token.lower()) if token.lower() in local_parameters else token.upper()
 
+        constructed = ARRAY_CONSTRUCTOR.search(text)
+        if constructed and constructed.span() == (0, len(text)):
+            return array_of(constructed)
+        if constructed:
+            # A constant expression over a constructor (CLUBB's saturation
+            # and pdf_closure: ``100._core_rknd * (/ 6.09868993_core_rknd,
+            # ... /)``): each constructor as an array where it stands, the
+            # rest through the token pass; NumPy is elementwise as Fortran
+            # is. A call around one (``reshape((/.../), ...)``) still goes
+            # to the parse, by the uppercased-call rule.
+            pieces: list[str] = []
+            position = 0
+            for found in ARRAY_CONSTRUCTOR.finditer(text):
+                pieces.append(IDENTIFIER.sub(case_of, strip_kind(text[position : found.start()])))
+                pieces.append(array_of(found))
+                position = found.end()
+            pieces.append(IDENTIFIER.sub(case_of, strip_kind(text[position:])))
+            return "".join(pieces)
+
         return IDENTIFIER.sub(case_of, strip_kind(text))
 
     @staticmethod
diff --git a/tests/test_flatten.py b/tests/test_flatten.py
index f0c451d..a7f079c 100644
--- a/tests/test_flatten.py
+++ b/tests/test_flatten.py
@@ -706,3 +706,76 @@ def test_the_recorder_writes_the_plain_out_dummies_too(tmp_path: Path) -> None:
     recorder = recorder_module("solver_mod", [plan])
     assert "call rec_r1(u_apply, 'INPUT', 'x', trim(dims), reshape(x, (/size(x)/)))" in recorder
     assert "call rec_r1(u_apply, 'OUTPUT', 'x', trim(dims), reshape(x, (/size(x)/)))" in recorder
+
+
+SPONGE_STATE = """\
+module sponge_mod
+  implicit none
+  private
+  public :: profile_type, damp, sponge_profile, init_profile
+  type profile_type
+    real(8), allocatable :: tau(:)
+    integer :: n_sponge = 0
+  end type profile_type
+  type(profile_type), public :: sponge_profile
+contains
+  subroutine init_profile( nz, prof )
+    integer, intent(in) :: nz
+    type(profile_type), intent(inout) :: prof
+    allocate( prof%tau(1:nz) )
+    prof%tau = 1.0d0
+    prof%n_sponge = nz
+  end subroutine init_profile
+  function damp( nzt, x, prof ) result( damped )
+    integer, intent(in) :: nzt
+    real(8), intent(in) :: x(nzt)
+    type(profile_type), intent(in) :: prof
+    real(8) :: damped(nzt)
+    if ( allocated( prof%tau ) ) then
+      damped = x * prof%tau(1:nzt)
+    else
+      damped = x
+    end if
+  end function damp
+end module sponge_mod
+"""
+
+HANDS_STATE_TO_FUNCTION = """\
+module advance_mod
+  use coefs_mod, only: coefs_type
+  use sponge_mod, only: damp, sponge_profile
+  implicit none
+  private
+  public :: advance
+contains
+  subroutine advance( nzt, ngrdcol, c, x )
+    integer, intent(in) :: nzt, ngrdcol
+    type(coefs_type), intent(in) :: c
+    real(8), intent(inout), dimension(ngrdcol, nzt) :: x
+    integer :: i
+    do i = 1, ngrdcol
+      x(i, :) = damp( nzt, x(i, :), sponge_profile ) * c%coef(i, 1:nzt)
+    end do
+  end subroutine advance
+end module advance_mod
+"""
+
+
+def test_module_state_handed_whole_to_a_function_is_carried(tmp_path: Path) -> None:
+    """CLUBB's advance_xm_wpxp passes sponge_layer_damping's profile object
+    to its sponge_damp_xm function, which reads ``profile%tau_sponge_damp``.
+    The walk followed only the caller's own dummies into callees, so the
+    profile was left to the module and the replay found it unallocated."""
+    (tmp_path / "coefs_mod.f90").write_text(COEFS)
+    (tmp_path / "sponge_mod.f90").write_text(SPONGE_STATE)
+    (tmp_path / "advance_mod.f90").write_text(HANDS_STATE_TO_FUNCTION)
+    frontend = FortranFrontend(flatten={"patch_count": "ngrdcol", "bounds_pattern": r"^ngrdcol$"})
+    unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:advance_mod")
+    facts = frontend.analyze(unit, tmp_path)
+    (plan,) = plans_for(facts, tmp_path, CLUBB_CONVENTIONS)
+    assert plan.usable, plan.unsupported
+    by_name = {o.name: o for o in plan.objects}
+    assert by_name["sponge_profile"].kind == "state"
+    assert by_name["sponge_profile"].module == "sponge_mod"
+    assert [c.name for c in by_name["sponge_profile"].components] == ["tau"]
+    assert "sponge_mod%sponge_profile" not in plan.left_to_module
diff --git a/tests/test_numpy_translate.py b/tests/test_numpy_translate.py
index 154269f..52451b1 100644
--- a/tests/test_numpy_translate.py
+++ b/tests/test_numpy_translate.py
@@ -487,3 +487,59 @@ def test_an_optional_handed_on_carries_its_own_presence(tmp_path: Path) -> None:
     finally:
         sys.path.remove(str(out))
         sys.modules.pop("relay_mod_numpy", None)
+
+
+SCALED_CONSTRUCTOR = """
+module fit_mod
+  implicit none
+  integer, parameter :: r8 = selected_real_kind(15)
+contains
+  subroutine polynomial( n, x, y )
+    integer, intent(in) :: n
+    real(r8), intent(in) :: x(n)
+    real(r8), intent(out) :: y(n)
+    real(r8), dimension(3), parameter :: &
+      a = 100._r8 * (/ 6.09868993_r8, 0.499320233_r8, 0.184672631E-01_r8 /)
+    real(r8), dimension(3), parameter :: b = (/ 1._r8, 2._r8, 3._r8 /) / 2._r8
+    integer :: i
+    do i = 1, n
+      y(i) = a(1) + a(2) * x(i) + a(3) * x(i)**2 + b(3)
+    end do
+  end subroutine polynomial
+end module fit_mod
+"""
+
+
+def test_a_constant_expression_over_an_array_constructor_is_a_value(tmp_path: Path) -> None:
+    """CLUBB's saturation and pdf_closure (#26): ``100._core_rknd * (/ ... /)``
+    as a local parameter. The token pass rendered a bare constructor and
+    handed anything around one to the parser, whose literals were never
+    hoisted; the whole subprogram was a NotImplementedError."""
+    import importlib
+    import sys
+
+    (tmp_path / "fit_mod.f90").write_text(SCALED_CONSTRUCTOR)
+    frontend = FortranFrontend()
+    unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:fit_mod")
+    facts = frontend.analyze(unit, tmp_path)
+    candidate = NumpyTranslation().apply(unit, facts, {"root": tmp_path})
+    assert not candidate.deferred
+    out = tmp_path / "emitted"
+    out.mkdir()
+    for path, content in candidate.files.items():
+        (out / path.name).write_bytes(content)
+    text = (out / "fit_mod_numpy.py").read_text()
+    assert "a = 100. * np.array([6.09868993, 0.499320233, 0.184672631E-01])" in text
+    assert "b = np.array([1., 2., 3.]) / 2." in text
+    sys.path.insert(0, str(out))
+    try:
+        module = importlib.import_module("fit_mod_numpy")
+        import numpy as np
+
+        y = module.polynomial(2, np.array([0.0, 1.0]))
+        a = 100.0 * np.array([6.09868993, 0.499320233, 0.184672631e-01])
+        # To rounding: the point is the constructor, not the summation order.
+        assert np.allclose(y, [a[0] + 1.5, a[0] + a[1] + a[2] + 1.5], rtol=0, atol=1e-9)
+    finally:
+        sys.path.remove(str(out))
+        sys.modules.pop("fit_mod_numpy", None)

From 5461d36bfca97441925837704bb52bf248e155ae Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 13:09:34 -0600
Subject: [PATCH 22/73] A subroutine with no OUT argument returns zero values,
 not one

CLUBB's finalize_tau_sponge_damp_api deallocates a component and returns;
its adapter returns None, which the gate counted as one value against
zero out-intent arguments.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/verify/bitexact.py | 14 +++++++--
 tests/test_f2py_oracle.py     | 55 +++++++++++++++++++++++++++++++++++
 2 files changed, 67 insertions(+), 2 deletions(-)

diff --git a/src/recast/verify/bitexact.py b/src/recast/verify/bitexact.py
index d00f118..de13971 100644
--- a/src/recast/verify/bitexact.py
+++ b/src/recast/verify/bitexact.py
@@ -51,6 +51,16 @@
 SUPPORTED_DTYPES = frozenset({"float32", "float64", "int32", "int64", "bool"})
 
 
+def _returned(translated_out: Any) -> list[Any]:
+    """The values a candidate call handed back, as a list: a tuple's items,
+    one bare value, or none at all -- a subroutine with no OUT argument
+    returns ``None`` (CLUBB's finalize_tau_sponge_damp_api deallocates and
+    returns), and that is zero values, not one."""
+    if translated_out is None:
+        return []
+    return list(translated_out) if isinstance(translated_out, tuple) else [translated_out]
+
+
 def _extent(dim: dict[str, Any], dims: dict[str, int]) -> int:
     """An axis's extent: ``ub - lb + 1`` when a lower bound is declared
     (CLUBB's ``lhs(-2:2, ...)`` has five rows, not two), ``ub`` otherwise."""
@@ -1060,7 +1070,7 @@ def _paired_outputs(
             # therefore by exact name on both sides. Every required output was
             # preflighted before the candidate call; keep the same check here
             # as a fail-closed local invariant for direct callers.
-            mine = list(translated_out) if isinstance(translated_out, tuple) else [translated_out]
+            mine = _returned(translated_out)
             names = (
                 [sub.get("result") or "result"]
                 if sub["kind"] == "function"
@@ -1116,7 +1126,7 @@ def _paired_outputs(
                 for a in outs_required
             ]
 
-        ours = list(translated_out) if isinstance(translated_out, tuple) else [translated_out]
+        ours = _returned(translated_out)
         if len(ours) != len(outs_all):
             return (
                 f"candidate returned {len(ours)} value(s) for "
diff --git a/tests/test_f2py_oracle.py b/tests/test_f2py_oracle.py
index f6a0d6f..8ac3e5c 100644
--- a/tests/test_f2py_oracle.py
+++ b/tests/test_f2py_oracle.py
@@ -718,6 +718,61 @@ def two_outputs():
     assert "partial output evidence is not a pass" in verdict.detail
 
 
+def test_a_recorded_subroutine_with_no_output_returns_nothing(tmp_path: Path) -> None:
+    """CLUBB's finalize_tau_sponge_damp_api deallocates a component and
+    returns: no OUT argument, so its adapter returns ``None``. The gate
+    counted that as one value against zero out-intent arguments."""
+    emitted = b"""\
+_SIGNATURES = {
+    "release": {
+        "kind": "subroutine",
+        "args": [
+            {"name": "n", "dtype": "int32", "intent": "IN", "optional": False},
+        ],
+        "result": None,
+        "result_dtype": None,
+    }
+}
+
+def release(n):
+    return None
+"""
+    candidate = Candidate(
+        unit="fortran:no_output",
+        transform="translate.numpy",
+        files={Path("no_output_numpy.py"): emitted},
+    )
+    ref = OracleRef(
+        unit=candidate.unit,
+        oracle="dump-replay",
+        key="k",
+        handle={
+            "module": None,
+            "input_source": "recorded",
+            "return_convention": "recorded",
+            "samples": [
+                {
+                    "subprogram": "release",
+                    "source": "release.txt",
+                    "inputs": {"n": 3},
+                    "outputs": {},
+                }
+            ],
+        },
+    )
+    verdict = BitexactVerifier().verify(
+        Unit(uid=candidate.unit, kind="module"),
+        candidate,
+        ref,
+        tmp_path / "work",
+        LocalExecutor(),
+        {},
+    )
+    # Nothing to compare is still not a pass -- but for the right reason.
+    assert "returned 1 value(s)" not in verdict.detail
+    assert "zero numerical points" in verdict.detail
+
+
 # --- the whole spine, against a real compiler --------------------------------
 
 SOURCE = """\

From 4f642504fadec1d6d35035796c70ab27b9a2c4c2 Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 13:13:37 -0600
Subject: [PATCH 23/73] A declared-ungated subprogram is not compared on a
 recording either

The declaration carries the reason -- CLUBB's sponge initializer leaves
the levels below the layer undefined on both sides -- and the replay
reported it, then compared heap contents against np.empty anyway.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/verify/bitexact.py | 14 +++++----
 tests/test_f2py_oracle.py     | 55 +++++++++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+), 5 deletions(-)

diff --git a/src/recast/verify/bitexact.py b/src/recast/verify/bitexact.py
index de13971..cfcca0d 100644
--- a/src/recast/verify/bitexact.py
+++ b/src/recast/verify/bitexact.py
@@ -270,6 +270,12 @@ def generable(name: str) -> bool:
                 if a["intent"] != "OUT"
             )
 
+        # One the operator declared ungated is not compared: the declaration
+        # says the reference cannot be held -- on generated inputs (CLUBB's
+        # rcm_sat_adj iterates and error-stops on them) or on a recording
+        # (its sponge initializer leaves the levels below the layer
+        # undefined on both sides). The reason is reported beside the verdict.
+        declared_ungated = set(config.get("ungated") or {})
         if recorded:
             # A recording names what it is a recording of, so the set to
             # compare is the set that was captured -- not every subprogram the
@@ -281,15 +287,13 @@ def generable(name: str) -> bool:
                 by_subprogram.setdefault(str(sample.get("subprogram", "")), []).append(sample)
             offered = sorted(by_subprogram)
             wanted = config.get("subprograms") or [
-                name for name in offered if name in table and judged(name)
+                name
+                for name in offered
+                if name in table and judged(name) and name not in declared_ungated
             ]
             skipped = sorted(set(offered) - set(wanted))
         else:
             by_subprogram = {}
-            # One the operator declared ungated is not sampled either: the
-            # declaration says the reference cannot be held on generated
-            # inputs (CLUBB's rcm_sat_adj iterates and error-stops on them).
-            declared_ungated = set(config.get("ungated") or {})
             wanted = config.get("subprograms") or [
                 name
                 for name in wrappers
diff --git a/tests/test_f2py_oracle.py b/tests/test_f2py_oracle.py
index 8ac3e5c..7cd5afe 100644
--- a/tests/test_f2py_oracle.py
+++ b/tests/test_f2py_oracle.py
@@ -773,6 +773,61 @@ def release(n):
     assert "zero numerical points" in verdict.detail
 
 
+def test_a_declared_ungated_subprogram_is_not_compared_on_a_recording(tmp_path: Path) -> None:
+    """CLUBB's sponge initializer leaves the levels below the layer undefined
+    on both sides; the operator's declaration says so, with the reason, and
+    the replay reported it -- then compared the heap against np.empty anyway."""
+    emitted = b"""\
+_SIGNATURES = {
+    "fill": {
+        "kind": "subroutine",
+        "args": [
+            {"name": "n", "dtype": "int32", "intent": "IN", "optional": False},
+            {"name": "y", "dtype": "int32", "intent": "OUT", "optional": False},
+        ],
+        "result": None,
+        "result_dtype": None,
+    }
+}
+
+def fill(n):
+    return 2
+"""
+    candidate = Candidate(
+        unit="fortran:undefined_tail",
+        transform="translate.numpy",
+        files={Path("undefined_tail_numpy.py"): emitted},
+    )
+    ref = OracleRef(
+        unit=candidate.unit,
+        oracle="dump-replay",
+        key="k",
+        handle={
+            "module": None,
+            "input_source": "recorded",
+            "return_convention": "recorded",
+            "samples": [
+                {
+                    "subprogram": "fill",
+                    "source": "fill.txt",
+                    "inputs": {"n": 3},
+                    "outputs": {"y": 1},
+                }
+            ],
+        },
+    )
+    verdict = BitexactVerifier().verify(
+        Unit(uid=candidate.unit, kind="module"),
+        candidate,
+        ref,
+        tmp_path / "work",
+        LocalExecutor(),
+        {"ungated": {"fill": "the tail is undefined on both sides"}},
+    )
+    assert "differ" not in verdict.detail
+    assert "fill (the tail is undefined on both sides)" in verdict.detail
+
+
 # --- the whole spine, against a real compiler --------------------------------
 
 SOURCE = """\

From c34e6c901953895d98c73981a66927fa7f2d57ea Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 13:58:32 -0600
Subject: [PATCH 24/73] Re-record the corpus baseline on this branch

No unit regresses; numfor's genrand goes from 7 to 3 disagreeing blocks and
slsqp from 12 to 8 (the component-name and handed-on-optional rules).

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 corpus/baseline.json | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/corpus/baseline.json b/corpus/baseline.json
index 1912df3..e22272e 100644
--- a/corpus/baseline.json
+++ b/corpus/baseline.json
@@ -1390,11 +1390,11 @@
      "stopped_by": "static.rwset",
      "verdicts": {
       "static.rwset": {
-       "detail": "7/187 blocks disagree: genrand_encode/B002, genrand_encode/B003, genrand_decode/B002, genrand_load_state/B003, genrand_load_state/B004 (+2 more)",
+       "detail": "3/187 blocks disagree: genrand_encode/B002, genrand_encode/B003, genrand_decode/B002",
        "metrics": {
         "blocks_checked": 187,
         "blocks_deferred": 3,
-        "blocks_matched": 180,
+        "blocks_matched": 184,
         "blocks_waived": 0
        },
        "passed": false
@@ -2108,11 +2108,11 @@
      "stopped_by": "static.rwset",
      "verdicts": {
       "static.rwset": {
-       "detail": "12/111 blocks disagree: slsqp/B019, slsqpb/B002, slsqpb/B003, slsqpb/B006, reset_bfgs_matrix/B002 (+7 more)",
+       "detail": "8/111 blocks disagree: slsqpb/B002, slsqpb/B003, slsqpb/B006, reset_bfgs_matrix/B002, lsq/B026 (+3 more)",
        "metrics": {
         "blocks_checked": 111,
         "blocks_deferred": 3,
-        "blocks_matched": 99,
+        "blocks_matched": 103,
         "blocks_waived": 0
        },
        "passed": false

From f254155a507b9f0b51ca5576d2e08d7608c57a2f Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 14:20:57 -0600
Subject: [PATCH 25/73] Re-record the corpus baseline after the merge with main

Against main's own baseline no unit regresses; slsqp goes from 15 to 7
disagreeing blocks of the same 106 checked.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 corpus/baseline.json | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/corpus/baseline.json b/corpus/baseline.json
index df75b0f..70891b3 100644
--- a/corpus/baseline.json
+++ b/corpus/baseline.json
@@ -1933,7 +1933,9 @@
           "max_rel": 0.0,
           "max_ulp": 0,
           "nan_mismatch": 0,
-          "points": 10
+          "points": 10,
+          "redrawn": 0,
+          "reshaped": 0
          },
          "searchsorted_idp": {
           "error": "candidate raised: NameError: name 'SMALL' is not defined"
@@ -2370,11 +2372,11 @@
      "stopped_by": "static.rwset",
      "verdicts": {
       "static.rwset": {
-       "detail": "8/111 blocks disagree: slsqpb/B002, slsqpb/B003, slsqpb/B006, reset_bfgs_matrix/B002, lsq/B026 (+3 more)",
+       "detail": "7/106 blocks disagree: slsqpb/B002, slsqpb/B003, slsqpb/B006, reset_bfgs_matrix/B002, lsq/B026 (+2 more)",
        "metrics": {
-        "blocks_checked": 111,
-        "blocks_deferred": 3,
-        "blocks_matched": 103,
+        "blocks_checked": 106,
+        "blocks_deferred": 8,
+        "blocks_matched": 99,
         "blocks_waived": 0
        },
        "passed": false

From 59e0fe84006a65b8989073f9fe090ad9b0692a49 Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 14:25:01 -0600
Subject: [PATCH 26/73] A probe spans blank and comment lines inside a
 continued call

cpp leaves blank lines where an #ifdef stood inside CLUBB's
advance_clubb_core call; the probe took one for the end of the statement.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/oracle/record.py | 14 +++++++++++++-
 tests/test_flatten.py       | 38 +++++++++++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/src/recast/oracle/record.py b/src/recast/oracle/record.py
index 35f3b0c..dacad47 100644
--- a/src/recast/oracle/record.py
+++ b/src/recast/oracle/record.py
@@ -353,6 +353,15 @@ def _continues(line: str) -> bool:
     return _strip_comment(line).rstrip().endswith("&")
 
 
+def _last_code(span: list[str]) -> str:
+    """The last line of ``span`` that carries code: blank and comment-only
+    lines do not end a continued statement."""
+    for line in reversed(span):
+        if _strip_comment(line).strip():
+            return line
+    return span[-1]
+
+
 def _split_actuals(text: str) -> list[str]:
     """Split an actual-argument list on top-level commas only."""
     parts: list[str] = []
@@ -398,8 +407,11 @@ def probe_tree(
         i = 0
         while i < len(lines):
             # A call may continue over several lines: gather the statement.
+            # A blank or comment-only line between continuations (what cpp
+            # leaves of an ``#ifdef`` inside CLUBB's advance_clubb_core call)
+            # is part of the statement, not its end.
             span = [lines[i]]
-            while _continues(span[-1]) and i + len(span) < len(lines):
+            while _continues(_last_code(span)) and i + len(span) < len(lines):
                 span.append(lines[i + len(span)])
             # ``a, & ! In`` leaves a space between the ampersand and the
             # comment it trailed: strip blanks before the ampersands.
diff --git a/tests/test_flatten.py b/tests/test_flatten.py
index a7f079c..bac823a 100644
--- a/tests/test_flatten.py
+++ b/tests/test_flatten.py
@@ -779,3 +779,41 @@ def test_module_state_handed_whole_to_a_function_is_carried(tmp_path: Path) -> N
     assert by_name["sponge_profile"].module == "sponge_mod"
     assert [c.name for c in by_name["sponge_profile"].components] == ["tau"]
     assert "sponge_mod%sponge_profile" not in plan.left_to_module
+
+
+def test_a_probe_spans_a_blank_line_inside_a_continued_call(tmp_path: Path) -> None:
+    """cpp leaves blank lines where an ``#ifdef`` stood inside CLUBB's
+    advance_clubb_core call; the probe took the blank line for the end of
+    the statement and found no call site to bracket."""
+    from recast.oracle.record import probe_tree
+
+    (tmp_path / "coefs_mod.f90").write_text(COEFS)
+    (tmp_path / "solver_mod.f90").write_text(USES_COEFS)
+    (tmp_path / "caller_mod.f90").write_text(
+        "module caller_mod\n"
+        "  use coefs_mod, only: coefs_type\n"
+        "  use solver_mod, only: apply\n"
+        "  implicit none\n"
+        "contains\n"
+        "  subroutine run( nzt, ngrdcol, c, x )\n"
+        "    integer, intent(in) :: nzt, ngrdcol\n"
+        "    type(coefs_type), intent(in) :: c\n"
+        "    real(8), intent(inout), dimension(ngrdcol, nzt) :: x\n"
+        "    call apply( nzt, ngrdcol, & ! in\n"
+        "\n"
+        "                ! the object\n"
+        "                c, &\n"
+        "\n"
+        "                x )\n"
+        "  end subroutine run\n"
+        "end module caller_mod\n"
+    )
+    frontend = FortranFrontend(flatten={"patch_count": "ngrdcol", "bounds_pattern": r"^ngrdcol$"})
+    unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:solver_mod")
+    facts = frontend.analyze(unit, tmp_path)
+    (plan,) = plans_for(facts, tmp_path, CLUBB_CONVENTIONS)
+    sites = probe_tree(tmp_path, tmp_path / "probed", {"solver_mod": [plan]})
+    assert sites == {"apply": 1}
+    probed = (tmp_path / "probed" / "caller_mod.f90").read_text()
+    assert "call rec_apply(0, nzt, ngrdcol, c, x)" in probed
+    assert "call rec_apply(1, nzt, ngrdcol, c, x)" in probed

From cc91f0ac7ed97db4bf4a7403ddc26305f4b13dde Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 14:32:48 -0600
Subject: [PATCH 27/73] A generic function reference is ranked by the specific
 it dispatches to

CLUBB's advance_clubb_core passes thlm2T_in_K_api( ... ) -- a generic
over a scalar and a 2-D specific -- as an actual to sat_mixrat_liq_api.
Ranked scalar without looking, the outer generic matched no specific, its
block was deferred, and a deferred block takes the whole subprogram out of
the gate. A non-elemental function's rank is its declared result's.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/fortran/semantics.py | 18 +++++--
 tests/test_numpy_translate.py   | 88 +++++++++++++++++++++++++++++++++
 2 files changed, 102 insertions(+), 4 deletions(-)

diff --git a/src/recast/fortran/semantics.py b/src/recast/fortran/semantics.py
index 4bdae54..f434f2c 100644
--- a/src/recast/fortran/semantics.py
+++ b/src/recast/fortran/semantics.py
@@ -442,8 +442,17 @@ def _reference_rank(self, node: Any) -> int:
             return 0
         if name in self.procedures:
             return self._call_rank(self.procedures[name], items)
-        if name in self.companion_generics:
-            return 0  # the overload decides, and dispatch is a separate question
+        if name in self.companion_generics or name in self.generics:
+            # The overload decides: CLUBB's ``sat_mixrat_liq_api( ...,
+            # thlm2T_in_K_api( nzt, ngrdcol, thlm, exner, rcm ), ... )`` is
+            # an actual whose rank is the inner generic's 2-D specific's
+            # result rank, and the outer generic's dispatch needs it.
+            try:
+                specific = self.dispatch(name, items)
+            except AmbiguousDispatch:
+                return 0
+            record = self.procedures.get(specific)
+            return self._call_rank(record, items) if record is not None else 0
         if name in TRANSFORMATIONAL or name in STATE_QUERY - {"merge"}:
             return 0
         if name in ELEMENTAL or name == "merge":
@@ -455,10 +464,11 @@ def _reference_rank(self, node: Any) -> int:
         return sum(1 for s in items if isinstance(s, f03.Subscript_Triplet))
 
     def _call_rank(self, record: dict[str, Any], items: list[Any]) -> int:
-        """An ELEMENTAL function broadcasts; anything else returns its result."""
+        """An ELEMENTAL function broadcasts; anything else returns its result
+        -- an array where the result is declared with dimensions."""
         if any("ELEMENTAL" in str(p).upper() for p in (record.get("prefixes") or [])):
             return self._broadcast_rank(items)
-        return 0
+        return len(record.get("result_dims") or [])
 
     def _broadcast_rank(self, items: list[Any]) -> int:
         return max(
diff --git a/tests/test_numpy_translate.py b/tests/test_numpy_translate.py
index ec2110f..c39757b 100644
--- a/tests/test_numpy_translate.py
+++ b/tests/test_numpy_translate.py
@@ -617,3 +617,91 @@ def test_the_callback_translation_passes_the_dataflow_gate(callback_candidate) -
     verdict = ReadWriteSetVerifier().check(unit, candidate, Path("."), LocalExecutor(), {})
     assert verdict.confidence is Confidence.SAMPLED, verdict.detail
     assert verdict.metrics["blocks_matched"] == verdict.metrics["blocks_checked"] > 0
+
+
+NESTED_GENERICS = """
+module sat_mod
+  implicit none
+  private
+  public :: t_api, rsat_api
+  interface t_api
+    module procedure t_k, t_2d
+  end interface t_api
+  interface rsat_api
+    module procedure rsat_k, rsat_2d
+  end interface rsat_api
+contains
+  function t_k( thl ) result( t )
+    real, intent(in) :: thl
+    real :: t
+    t = thl + 1.0
+  end function t_k
+  function t_2d( n, thl ) result( t )
+    integer, intent(in) :: n
+    real, intent(in) :: thl(n, 2)
+    real :: t(n, 2)
+    t = thl + 1.0
+  end function t_2d
+  function rsat_k( p, t ) result( r )
+    real, intent(in) :: p, t
+    real :: r
+    r = t / p
+  end function rsat_k
+  function rsat_2d( n, p, t ) result( r )
+    integer, intent(in) :: n
+    real, intent(in) :: p(n, 2), t(n, 2)
+    real :: r(n, 2)
+    r = t / p
+  end function rsat_2d
+end module sat_mod
+"""
+
+CALLS_NESTED_GENERICS = """
+module core_mod
+  use sat_mod, only: t_api, rsat_api
+  implicit none
+contains
+  subroutine step( n, p, thl, rsat )
+    integer, intent(in) :: n
+    real, intent(in) :: p(n, 2), thl(n, 2)
+    real, intent(out) :: rsat(n, 2)
+    rsat = rsat_api( n, p, t_api( n, thl ) )
+  end subroutine step
+end module core_mod
+"""
+
+
+def test_a_generic_whose_actual_is_another_generics_result_is_dispatched(tmp_path: Path) -> None:
+    """CLUBB's advance_clubb_core: ``sat_mixrat_liq_api( ..., thlm2T_in_K_api(
+    ... ), ... )``. The inner generic's result was ranked scalar without
+    looking, so the outer one matched no specific and its block was deferred
+    -- and a deferred block takes the whole subprogram out of the gate."""
+    import importlib
+    import sys
+
+    (tmp_path / "sat_mod.f90").write_text(NESTED_GENERICS)
+    (tmp_path / "core_mod.f90").write_text(CALLS_NESTED_GENERICS)
+    frontend = FortranFrontend()
+    unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:core_mod")
+    facts = frontend.analyze(unit, tmp_path)
+    from recast.transform.numpy.tree import TreeTranslation
+
+    candidate = TreeTranslation().apply(unit, facts, {"root": str(tmp_path)})
+    assert not candidate.deferred, candidate.deferred
+    out = tmp_path / "emitted"
+    out.mkdir()
+    for path, content in candidate.files.items():
+        (out / path.name).write_bytes(content)
+    text = (out / "core_mod_numpy.py").read_text()
+    assert "rsat_2d(" in text and "t_2d(" in text
+    sys.path.insert(0, str(out))
+    try:
+        module = importlib.import_module("core_mod_numpy")
+        import numpy as np
+
+        p = np.full((2, 2), 2.0, dtype=np.float32, order="F")
+        thl = np.ones((2, 2), dtype=np.float32, order="F")
+        assert module.step(2, p, thl).tolist() == [[1.0, 1.0], [1.0, 1.0]]
+    finally:
+        sys.path.remove(str(out))
+        sys.modules.pop("core_mod_numpy", None)

From 59faf7af7f84fcf5212ae5eba7360be3e76909d6 Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 14:43:15 -0600
Subject: [PATCH 28/73] JAX lowering: a CYCLE folds into the branch structure;
 an EXIT or a leftover CYCLE is delegated

CLUBB's interpolators cycle from inside an IF in a DO loop. The lowering
passed the continue through into a lax.cond branch -- a SyntaxError that
took the whole emitted module down. The loop body's continuation now moves
into the branches that do not cycle, before the fori_loop lowering; what
that cannot fold, and every break, delegates the subprogram to the host.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/jax/backend.py | 44 +++++++++++++++++-
 tests/test_jax_transform.py         | 71 +++++++++++++++++++++++++++++
 2 files changed, 114 insertions(+), 1 deletion(-)

diff --git a/src/recast/transform/jax/backend.py b/src/recast/transform/jax/backend.py
index 3232802..054dfa5 100644
--- a/src/recast/transform/jax/backend.py
+++ b/src/recast/transform/jax/backend.py
@@ -291,6 +291,42 @@ def visit_Call(self, node):
 # ------------------------------------------------------- statement lowering
 
 
+def _has_cycle(stmts) -> bool:
+    """Whether a ``continue`` of *this* loop sits in ``stmts`` -- at the top
+    level or under an ``if``; one inside a nested ``for`` is that loop's."""
+    for s in stmts:
+        if isinstance(s, ast.Continue):
+            return True
+        if isinstance(s, ast.If) and (_has_cycle(s.body) or _has_cycle(s.orelse)):
+            return True
+    return False
+
+
+def _cycle_to_else(stmts, rest):
+    """``stmts`` followed by ``rest``, with every ``continue`` folded away.
+
+    Fortran's ``if ( ... ) then ... cycle end if`` followed by the rest of
+    the loop body (CLUBB's interpolators) is ``if c: A else: R`` -- the
+    continuation of the loop body moves into the branches that do not
+    cycle. Done on the Python AST before the fori_loop lowering, which has
+    no place for a ``continue``: a ``lax.cond`` branch is a function.
+    """
+    out = []
+    for at, s in enumerate(stmts):
+        if isinstance(s, ast.Continue):
+            return out  # what follows is never reached on this path
+        if isinstance(s, ast.If) and (_has_cycle(s.body) or _has_cycle(s.orelse)):
+            tail = [*stmts[at + 1 :], *rest]
+            folded = ast.If(
+                test=s.test,
+                body=_cycle_to_else(s.body, copy.deepcopy(tail)) or [ast.Pass()],
+                orelse=_cycle_to_else(s.orelse, copy.deepcopy(tail)),
+            )
+            return [*out, ast.copy_location(folded, s)]
+        out.append(s)
+    return [*out, *rest]
+
+
 def _assigned_names(stmts):
     """Names stored by Assign statements, first-assignment order
     (nested fori_loop/cond results arrive as Tuple targets; static
@@ -453,6 +489,12 @@ def lower_block(self, stmts, depth):
         for s in stmts:
             if isinstance(s, self.BANNED):
                 raise JaxQueue(f"unsupported stmt {type(s).__name__}")
+            if isinstance(s, ast.Continue | ast.Break):
+                # A CYCLE the loop pass could not fold into a branch, or an
+                # EXIT: neither has a place in a fori_loop body. Delegated,
+                # not emitted -- a ``continue`` inside a ``lax.cond`` branch
+                # is a SyntaxError that takes the whole module down.
+                raise JaxQueue(f"{type(s).__name__.lower()} inside a lowered loop")
             if isinstance(s, ast.Return) and depth > 0:
                 raise JaxQueue("return inside loop/branch body")
             if isinstance(s, ast.For):
@@ -566,7 +608,7 @@ def lower_for(self, s, depth):
         else:
             raise JaxQueue("malformed range")
 
-        body = self.lower_block(s.body, depth + 1)
+        body = self.lower_block(_cycle_to_else(s.body, []), depth + 1)
         settled = _trace_constant_stores(body)
         carried = [n for n in _assigned_names(body) if n != s.target.id and n not in settled]
         if not carried:
diff --git a/tests/test_jax_transform.py b/tests/test_jax_transform.py
index dae7354..f85cbf9 100644
--- a/tests/test_jax_transform.py
+++ b/tests/test_jax_transform.py
@@ -241,3 +241,74 @@ def test_a_constant_table_is_read_through_jnp() -> None:
     text = _ast.unparse(node)
     assert "jnp.asarray(MDAYLEAP)[m - 1]" in text
     assert "x[i]" in text  # a lowercase array is a traced value already
+
+
+CYCLES = """\
+module cycle_demo
+  implicit none
+  integer, parameter :: r8 = selected_real_kind(12)
+contains
+  subroutine clip_below(n, zlo, z, v, w)
+    integer,  intent(in)  :: n
+    real(r8), intent(in)  :: zlo
+    real(r8), intent(in)  :: z(n), v(n)
+    real(r8), intent(out) :: w(n)
+    integer :: k
+    do k = 1, n
+      if ( z(k) < zlo ) then
+        w(k) = 0.0_r8
+        cycle
+      end if
+      w(k) = v(k) * 2.0_r8
+    end do
+  end subroutine clip_below
+  subroutine first_above(n, zlo, z, kfound)
+    integer,  intent(in)  :: n
+    real(r8), intent(in)  :: zlo
+    real(r8), intent(in)  :: z(n)
+    integer,  intent(out) :: kfound
+    integer :: k
+    kfound = 0
+    do k = 1, n
+      if ( z(k) > zlo ) then
+        kfound = k
+        exit
+      end if
+    end do
+  end subroutine first_above
+end module cycle_demo
+"""
+
+
+def test_a_cycle_folds_into_the_branch_and_an_exit_is_delegated(tmp_path: Path) -> None:
+    """CLUBB's interpolators: ``if ( ... ) then ... cycle end if`` in a DO
+    loop. The lowering passed the ``continue`` through into a ``lax.cond``
+    branch -- a SyntaxError that took the whole emitted module down. Folded
+    into the branch structure it is a kernel; an EXIT has no fori_loop
+    shape and is delegated to the host, not emitted."""
+    import importlib
+    import sys
+
+    candidate = port(tmp_path, CYCLES, "cycle_demo")
+    assert candidate.notes["jax"]["kernels"] == ["clip_below"]
+    assert "first_above" in candidate.notes["jax"]["delegated"]
+    emitted = candidate.files[Path("cycle_demo_jax.py")].decode()
+    assert "continue" not in emitted.replace("continuation", "")
+    out = tmp_path / "emitted"
+    out.mkdir()
+    for path, content in candidate.files.items():
+        (out / path.name).write_bytes(content)
+    sys.path.insert(0, str(out))
+    try:
+        module = importlib.import_module("cycle_demo_jax")
+        import numpy as np
+
+        z = np.array([1.0, 2.0, 3.0, 4.0])
+        v = np.array([1.0, 1.0, 1.0, 1.0])
+        w = np.asarray(module.clip_below(4, 2.5, z, v))
+        assert w.tolist() == [0.0, 0.0, 2.0, 2.0]
+        assert int(module.first_above(4, 2.5, z)) == 3
+    finally:
+        sys.path.remove(str(out))
+        for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"):
+            sys.modules.pop(f"cycle_demo{suffix}", None)

From b2fba33a1f5100e8428fec42ab89af7d0c0c3ab3 Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 14:43:32 -0600
Subject: [PATCH 29/73] Annotate the folded statement list (mypy)

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/jax/backend.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/recast/transform/jax/backend.py b/src/recast/transform/jax/backend.py
index 054dfa5..6e9d210 100644
--- a/src/recast/transform/jax/backend.py
+++ b/src/recast/transform/jax/backend.py
@@ -311,7 +311,7 @@ def _cycle_to_else(stmts, rest):
     cycle. Done on the Python AST before the fori_loop lowering, which has
     no place for a ``continue``: a ``lax.cond`` branch is a function.
     """
-    out = []
+    out: list[ast.stmt] = []
     for at, s in enumerate(stmts):
         if isinstance(s, ast.Continue):
             return out  # what follows is never reached on this path

From bc1560d3f2c5b0f799805a0c74b5270c411b0a4e Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 14:46:46 -0600
Subject: [PATCH 30/73] JAX runtime: sqrt, ordered SUM, erf and erfc

CLUBB's clipping and PDF closure reach the kernels as _f_sqrt, _f_vsum and
_f_verf, which the NumPy runtime defines and the JAX one did not: a
NameError on the first call of every kernel of the unit.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/jax/runtime.py | 43 ++++++++++++++++++++++++++
 tests/test_jax_transform.py         | 47 +++++++++++++++++++++++++++++
 2 files changed, 90 insertions(+)

diff --git a/src/recast/transform/jax/runtime.py b/src/recast/transform/jax/runtime.py
index 05d3973..5468638 100644
--- a/src/recast/transform/jax/runtime.py
+++ b/src/recast/transform/jax/runtime.py
@@ -51,7 +51,11 @@
     "_f_tiny",
     "_f_trim",
     "_f_vceil",
+    "_f_sqrt",
     "_f_vdot",
+    "_f_verf",
+    "_f_verfc",
+    "_f_vsum",
     "_f_vexp",
     "_f_vfloor",
     "_f_vlog",
@@ -147,6 +151,45 @@ def _f_vpow(a, b):
     return jnp.asarray(a) ** b
 
 
+def _f_sqrt(x):
+    """Fortran SQRT: a NaN for a negative real, not an exception, and the
+    correctly rounded root otherwise -- what the NumPy shim does with
+    math.sqrt, and what jnp.sqrt does by itself."""
+    return jnp.sqrt(x)
+
+
+def _f_vsum(a, axis=None):
+    """Fortran SUM accumulates in element order; a sequential fori_loop
+    keeps the fold order the NumPy anchor uses (its ``_f_vsum``), so the
+    two sides differ by XLA's rounding alone and not by association."""
+    arr = jnp.asarray(a)
+    if axis is None:
+        flat = jnp.ravel(arr, order="F")
+
+        def body(i, s):
+            return s + flat[i]
+
+        return lax.fori_loop(0, flat.shape[0], body, jnp.zeros((), dtype=arr.dtype))
+    moved = jnp.moveaxis(arr, axis, 0)
+
+    def body_axis(i, s):
+        return s + moved[i]
+
+    return lax.fori_loop(0, moved.shape[0], body_axis, jnp.zeros(moved.shape[1:], dtype=arr.dtype))
+
+
+def _f_verf(x):
+    from jax.scipy.special import erf as _erf
+
+    return _erf(x)
+
+
+def _f_verfc(x):
+    from jax.scipy.special import erfc as _erfc
+
+    return _erfc(x)
+
+
 def _f_vdot(a, b):
     """Fortran DOT_PRODUCT accumulates in order; sequential fori_loop
     keeps the fold order (XLA may still contract the FMA)."""
diff --git a/tests/test_jax_transform.py b/tests/test_jax_transform.py
index f85cbf9..ace83f2 100644
--- a/tests/test_jax_transform.py
+++ b/tests/test_jax_transform.py
@@ -312,3 +312,50 @@ def test_a_cycle_folds_into_the_branch_and_an_exit_is_delegated(tmp_path: Path)
         sys.path.remove(str(out))
         for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"):
             sys.modules.pop(f"cycle_demo{suffix}", None)
+
+
+SHIMS = """\
+module shim_demo
+  implicit none
+  integer, parameter :: r8 = selected_real_kind(12)
+contains
+  subroutine norms(n, x, root, total, err)
+    integer,  intent(in)  :: n
+    real(r8), intent(in)  :: x(n)
+    real(r8), intent(out) :: root(n), total, err(n)
+    root = sqrt( x )
+    total = sum( x )
+    err = erf( x )
+  end subroutine norms
+end module shim_demo
+"""
+
+
+def test_the_jax_runtime_carries_sqrt_sum_and_erf(tmp_path: Path) -> None:
+    """CLUBB's clipping and PDF closure: ``sqrt``, ``sum`` and ``erf`` reach
+    the kernels as ``_f_sqrt``, ``_f_vsum`` and ``_f_verf``, which the NumPy
+    runtime defines and the JAX one did not -- a NameError at the first
+    call, on every kernel of the unit."""
+    import importlib
+    import sys
+
+    candidate = port(tmp_path, SHIMS, "shim_demo")
+    assert candidate.notes["jax"]["kernels"] == ["norms"]
+    out = tmp_path / "emitted"
+    out.mkdir()
+    for path, content in candidate.files.items():
+        (out / path.name).write_bytes(content)
+    sys.path.insert(0, str(out))
+    try:
+        module = importlib.import_module("shim_demo_jax")
+        import numpy as np
+
+        x = np.array([4.0, 9.0, -1.0])
+        root, total, err = (np.asarray(v) for v in module.norms(3, x))
+        assert root[:2].tolist() == [2.0, 3.0] and np.isnan(root[2])
+        assert total == 12.0
+        assert abs(err[0] - 0.9999999845827421) < 1e-12
+    finally:
+        sys.path.remove(str(out))
+        for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"):
+            sys.modules.pop(f"shim_demo{suffix}", None)

From 2d4cf2595d8dd5eec7c1f8821a3d687ec730a9e7 Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 14:46:59 -0600
Subject: [PATCH 31/73] Sort the JAX runtime's export list (ruff)

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/jax/runtime.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/recast/transform/jax/runtime.py b/src/recast/transform/jax/runtime.py
index 5468638..3a1532c 100644
--- a/src/recast/transform/jax/runtime.py
+++ b/src/recast/transform/jax/runtime.py
@@ -48,14 +48,13 @@
     "_f_modulo",
     "_f_nint",
     "_f_sign",
+    "_f_sqrt",
     "_f_tiny",
     "_f_trim",
     "_f_vceil",
-    "_f_sqrt",
     "_f_vdot",
     "_f_verf",
     "_f_verfc",
-    "_f_vsum",
     "_f_vexp",
     "_f_vfloor",
     "_f_vlog",
@@ -63,6 +62,7 @@
     "_f_vmax",
     "_f_vmin",
     "_f_vpow",
+    "_f_vsum",
     "_fstr_eq",
     "jax",
     "jnp",

From 428968d354c5bb9a01ea5ae46529167edf18d033 Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 14:49:16 -0600
Subject: [PATCH 32/73] JAX lowering: a loop over a static empty trip count is
 skipped, not traced

CLUBB's scalar tracers under sclr_dim = 0: a DO over the zero-extent axis
runs no iteration, but fori_loop traced the body once and JAX refuses any
index into a size-0 axis at trace time. Loops go through the runtime's
_f_fori, which returns the carry unchanged when the bounds are static and
empty and is lax.fori_loop otherwise.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/jax/backend.py | 12 +++-----
 src/recast/transform/jax/runtime.py | 14 +++++++++
 tests/test_jax_transform.py         | 47 +++++++++++++++++++++++++++++
 tests/test_jax_tree.py              |  2 +-
 4 files changed, 67 insertions(+), 8 deletions(-)

diff --git a/src/recast/transform/jax/backend.py b/src/recast/transform/jax/backend.py
index 6e9d210..cdee415 100644
--- a/src/recast/transform/jax/backend.py
+++ b/src/recast/transform/jax/backend.py
@@ -615,13 +615,11 @@ def lower_for(self, s, depth):
             raise JaxQueue("loop with no carried effects")
         self.n += 1
         fname = f"_body_{self.n}"
-        carry_call = ast.Call(
-            func=ast.Attribute(
-                value=ast.Name(id="lax", ctx=ast.Load()), attr="fori_loop", ctx=ast.Load()
-            ),
-            args=[],
-            keywords=[],
-        )
+        # Through the runtime's ``_f_fori``: a trip count that is static
+        # and empty (a loop over an array's zero-extent axis -- CLUBB's
+        # scalar tracers under sclr_dim = 0) is skipped rather than traced,
+        # because JAX refuses any index into a size-0 axis at trace time.
+        carry_call = ast.Call(func=ast.Name(id="_f_fori", ctx=ast.Load()), args=[], keywords=[])
         result = ast.Assign(
             targets=[ast.Tuple(elts=_names(carried, ast.Store), ctx=ast.Store())], value=carry_call
         )
diff --git a/src/recast/transform/jax/runtime.py b/src/recast/transform/jax/runtime.py
index 3a1532c..b162a65 100644
--- a/src/recast/transform/jax/runtime.py
+++ b/src/recast/transform/jax/runtime.py
@@ -39,6 +39,7 @@
     "_f_adjustl",
     "_f_dim",
     "_f_epsilon",
+    "_f_fori",
     "_f_huge",
     "_f_int_div",
     "_f_len_trim",
@@ -151,6 +152,19 @@ def _f_vpow(a, b):
     return jnp.asarray(a) ** b
 
 
+def _f_fori(lo, hi, body, init):
+    """``lax.fori_loop`` unless the trip count is static and empty.
+
+    A Fortran DO over an array's zero-extent axis (CLUBB's scalar tracers
+    under ``sclr_dim = 0``) runs no iteration; ``fori_loop`` would still
+    trace the body once, and JAX refuses any index into a size-0 axis at
+    trace time. A dynamic bound is left to ``fori_loop``.
+    """
+    if isinstance(lo, int) and isinstance(hi, int) and hi <= lo:
+        return init
+    return lax.fori_loop(lo, hi, body, init)
+
+
 def _f_sqrt(x):
     """Fortran SQRT: a NaN for a negative real, not an exception, and the
     correctly rounded root otherwise -- what the NumPy shim does with
diff --git a/tests/test_jax_transform.py b/tests/test_jax_transform.py
index ace83f2..d4471d6 100644
--- a/tests/test_jax_transform.py
+++ b/tests/test_jax_transform.py
@@ -359,3 +359,50 @@ def test_the_jax_runtime_carries_sqrt_sum_and_erf(tmp_path: Path) -> None:
         sys.path.remove(str(out))
         for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"):
             sys.modules.pop(f"shim_demo{suffix}", None)
+
+
+EMPTY_AXIS = """\
+module tracer_demo
+  implicit none
+  integer, parameter :: r8 = selected_real_kind(12)
+contains
+  subroutine scale_tracers(n, m, f, x, y)
+    integer,  intent(in)  :: n, m
+    real(r8), intent(in)  :: f
+    real(r8), intent(in)  :: x(n, m)
+    real(r8), intent(out) :: y(n, m)
+    integer :: s
+    do s = 1, m
+      y(:, s) = f * x(:, s)
+    end do
+  end subroutine scale_tracers
+end module tracer_demo
+"""
+
+
+def test_a_loop_over_a_zero_extent_axis_runs_no_iteration(tmp_path: Path) -> None:
+    """CLUBB's scalar tracers under ``sclr_dim = 0``: ``do sclr = 1, sclr_dim``
+    over ``(ngrdcol, nzm, 0)`` arrays. ``fori_loop`` traced the body once
+    and JAX refused the index into the size-0 axis."""
+    import importlib
+    import sys
+
+    candidate = port(tmp_path, EMPTY_AXIS, "tracer_demo")
+    assert candidate.notes["jax"]["kernels"] == ["scale_tracers"]
+    out = tmp_path / "emitted"
+    out.mkdir()
+    for path, content in candidate.files.items():
+        (out / path.name).write_bytes(content)
+    sys.path.insert(0, str(out))
+    try:
+        module = importlib.import_module("tracer_demo_jax")
+        import numpy as np
+
+        empty = np.zeros((3, 0), order="F")
+        assert np.asarray(module.scale_tracers(3, 0, 2.0, empty)).shape == (3, 0)
+        x = np.ones((3, 2), order="F")
+        assert np.asarray(module.scale_tracers(3, 2, 2.0, x)).tolist() == [[2.0, 2.0]] * 3
+    finally:
+        sys.path.remove(str(out))
+        for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"):
+            sys.modules.pop(f"tracer_demo{suffix}", None)
diff --git a/tests/test_jax_tree.py b/tests/test_jax_tree.py
index d8bb8d7..01d7c19 100644
--- a/tests/test_jax_tree.py
+++ b/tests/test_jax_tree.py
@@ -76,7 +76,7 @@ def test_the_port_emits_the_flat_function_as_a_kernel(tree: Path) -> None:
     candidate = TreeToJax(CONVENTIONS).apply(unit, facts, {"root": str(tree)})
     ported = candidate.files[Path("physics_mod_jax.py")].decode()
     assert "def _warm_flat_k_impl(" in ported
-    assert "lax.fori_loop" in ported and ".at[p - 1, ic - 1].set(" in ported
+    assert "_f_fori(" in ported and ".at[p - 1, ic - 1].set(" in ported
     assert "_JAX_KERNELS = ['reset_flat', 'warm_flat']" in ported
     # The originals take the object and stay host-delegated; the flat
     # signatures reach the ported module's table for the gate.

From b1afef89cab85cb13abd4177f408d94d82c7f6fc Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 14:51:57 -0600
Subject: [PATCH 33/73] Tolerance gate: a zero-extent output has no dominant
 value to weigh

CLUBB's scalar tracers under sclr_dim = 0 are recorded as (1, 88, 0)
outputs; the dominant-value mask took the maximum of an empty array and
the unit's verdict was a plugin exception.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/verify/bitexact.py |  4 +++
 tests/test_f2py_oracle.py     | 63 +++++++++++++++++++++++++++++++++++
 2 files changed, 67 insertions(+)

diff --git a/src/recast/verify/bitexact.py b/src/recast/verify/bitexact.py
index 0f26ddc..7b03cd1 100644
--- a/src/recast/verify/bitexact.py
+++ b/src/recast/verify/bitexact.py
@@ -1158,6 +1158,10 @@ def _dominance(
         if dominant_at is None:
             return None
         magnitude = np.abs(reference)
+        if magnitude.size == 0:
+            # A zero-extent output (CLUBB's scalar tracers under
+            # sclr_dim = 0): nothing to weigh, and no maximum to take.
+            return []
         if axis in ("all", None) or magnitude.ndim <= 1:
             scale = magnitude.max()
         else:
diff --git a/tests/test_f2py_oracle.py b/tests/test_f2py_oracle.py
index acbbbfb..96e2ecf 100644
--- a/tests/test_f2py_oracle.py
+++ b/tests/test_f2py_oracle.py
@@ -1026,6 +1026,69 @@ def fill(n):
     assert "fill (the tail is undefined on both sides)" in verdict.detail
 
 
+def test_a_zero_extent_output_has_no_dominant_value_to_weigh(tmp_path: Path) -> None:
+    """CLUBB's scalar tracers under ``sclr_dim = 0``: an output of shape
+    ``(1, 88, 0)`` on the recording. The tolerance gate's dominant-value
+    mask took the maximum of an empty array and the whole unit's verdict
+    was a plugin exception."""
+    import numpy as np
+
+    from recast.verify.tolerance import ToleranceVerifier
+
+    emitted = b"""\
+import numpy as np
+_SIGNATURES = {
+    "tracers": {
+        "kind": "subroutine",
+        "args": [
+            {"name": "n", "dtype": "int32", "intent": "IN", "optional": False},
+            {"name": "y", "dtype": "float64", "intent": "OUT", "optional": False,
+             "dims": [{"lb": "1", "ub": "n"}, {"lb": "1", "ub": "0"}]},
+            {"name": "z", "dtype": "float64", "intent": "OUT", "optional": False},
+        ],
+        "result": None,
+        "result_dtype": None,
+    }
+}
+
+def tracers(n):
+    return np.zeros((n, 0)), 2.0
+"""
+    candidate = Candidate(
+        unit="fortran:tracers",
+        transform="translate.numpy",
+        files={Path("tracers_numpy.py"): emitted},
+    )
+    ref = OracleRef(
+        unit=candidate.unit,
+        oracle="dump-replay",
+        key="k",
+        handle={
+            "module": None,
+            "input_source": "recorded",
+            "return_convention": "recorded",
+            "samples": [
+                {
+                    "subprogram": "tracers",
+                    "source": "tracers.txt",
+                    "inputs": {"n": 3},
+                    "outputs": {"y": np.zeros((3, 0)), "z": 2.0},
+                }
+            ],
+        },
+    )
+    verdict = ToleranceVerifier().verify(
+        Unit(uid=candidate.unit, kind="module"),
+        candidate,
+        ref,
+        tmp_path / "work",
+        LocalExecutor(),
+        {"module_suffix": "_numpy.py", "dominant_axis": "all", "rel_scale": "array"},
+    )
+    assert "exception" not in verdict.detail and "zero-size" not in verdict.detail
+    assert verdict.confidence is not Confidence.FAILED, verdict.detail
+
+
 # --- the whole spine, against a real compiler --------------------------------
 
 SOURCE = """\

From 956fa86fe83fcb1216c3455b4276ab92e4dff94f Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 15:07:53 -0600
Subject: [PATCH 34/73] A flat kernel spells the optional dummy the plan leaves
 out

CLUBB's grid interpolators take an optional zt_min; the plan drops it (the
adapter calls with it absent) and the kernel inlines a body that tests
zt_min is not None -- a NameError on every call. The absence is bound
at the top of the kernel: None, and want_<name> = False for an optional
OUT.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/jax/tree.py | 33 ++++++++++++++++
 tests/test_jax_transform.py      | 64 ++++++++++++++++++++++++++++++++
 2 files changed, 97 insertions(+)

diff --git a/src/recast/transform/jax/tree.py b/src/recast/transform/jax/tree.py
index 840a26f..14f6eca 100644
--- a/src/recast/transform/jax/tree.py
+++ b/src/recast/transform/jax/tree.py
@@ -897,6 +897,7 @@ def _specialize(self, name: str, call: ast.Call, source: dict[str, Any]) -> Flat
         if not body or not isinstance(body[-1], ast.Return):
             body.append(ast.Return(value=_tuple(_outputs(plan))))
         taken = [_py(a["name"]) for a in plan.flat_args if a["intent"] != "OUT"]
+        body = [*_absent_optionals(plan, taken), *body]
         flat = ast.FunctionDef(
             name=plan.name,
             args=ast.arguments(
@@ -2297,6 +2298,37 @@ def _static_names(args: list[dict[str, Any]]) -> frozenset[str]:
     )
 
 
+def _absent_optionals(plan: FlatPlan, taken: list[str]) -> list[ast.stmt]:
+    """Bindings for the optional dummies the flat signature leaves out.
+
+    The plan drops an optional dummy: the adapter calls the original with it
+    absent, so its NumPy default (``None``) stands in. The kernel inlines
+    the original's body, which still names it -- CLUBB's grid interpolators
+    test ``present(zt_min)`` -- so the absence is spelled at the top:
+    ``zt_min = None``, and ``want_zt_min = False`` for an optional OUT, the
+    anchor's presence sentinel. A trace-time ``if x is not None`` then skips
+    the branch the Fortran skipped.
+    """
+    out: list[ast.stmt] = []
+    for a in plan.subprogram.get("args") or ():
+        if not a.get("optional"):
+            continue
+        name = _py(a["name"])
+        if name in taken:
+            continue
+        out.append(
+            ast.Assign(targets=[ast.Name(id=name, ctx=ast.Store())], value=ast.Constant(value=None))
+        )
+        if a.get("intent") == "OUT":
+            out.append(
+                ast.Assign(
+                    targets=[ast.Name(id=f"want_{name}", ctx=ast.Store())],
+                    value=ast.Constant(value=False),
+                )
+            )
+    return out
+
+
 def flat_function(
     fn: ast.FunctionDef,
     plan: FlatPlan,
@@ -2325,6 +2357,7 @@ def flat_function(
     if not body or not isinstance(body[-1], ast.Return):
         body.append(ast.Return(value=_tuple(_outputs(plan))))
     taken = [_py(a["name"]) for a in plan.flat_args if a["intent"] != "OUT"]
+    body = [*_absent_optionals(plan, taken), *body]
     flat = ast.FunctionDef(
         name=plan.name,
         args=ast.arguments(
diff --git a/tests/test_jax_transform.py b/tests/test_jax_transform.py
index d4471d6..6dbe553 100644
--- a/tests/test_jax_transform.py
+++ b/tests/test_jax_transform.py
@@ -406,3 +406,67 @@ def test_a_loop_over_a_zero_extent_axis_runs_no_iteration(tmp_path: Path) -> Non
         sys.path.remove(str(out))
         for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"):
             sys.modules.pop(f"tracer_demo{suffix}", None)
+
+
+FLOORED = """\
+module floor_mod
+  implicit none
+  type coefs_type
+    real(8), allocatable :: coef(:, :)
+  end type coefs_type
+contains
+  subroutine init_coefs( nz, ngrdcol, c )
+    integer, intent(in) :: nz, ngrdcol
+    type(coefs_type), intent(inout) :: c
+    allocate( c%coef(1:ngrdcol, 1:nz) )
+    c%coef = 2.0d0
+  end subroutine init_coefs
+  subroutine apply( nzt, ngrdcol, c, x, floor )
+    integer, intent(in) :: nzt, ngrdcol
+    type(coefs_type), intent(in) :: c
+    real(8), intent(inout), dimension(ngrdcol, nzt) :: x
+    real(8), intent(in), optional :: floor
+    x = x * c%coef(:, 1:nzt)
+    if ( present( floor ) ) then
+      x = max( x, floor )
+    end if
+  end subroutine apply
+end module floor_mod
+"""
+
+
+def test_a_flat_kernel_spells_the_optional_dummy_the_plan_leaves_out(tmp_path: Path) -> None:
+    """CLUBB's grid interpolators take an optional ``zt_min``; the plan drops
+    it (the adapter calls with it absent) and the kernel inlines a body that
+    tests ``zt_min is not None`` -- a NameError on every call."""
+    import importlib
+    import sys
+
+    from recast.transform.jax.tree import TreeToJax
+    from recast.transform.numpy.tree import TreeConventions
+
+    (tmp_path / "floor_mod.f90").write_text(FLOORED)
+    frontend = FortranFrontend(flatten={"patch_count": "ngrdcol", "bounds_pattern": r"^ngrdcol$"})
+    unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:floor_mod")
+    facts = frontend.analyze(unit, tmp_path)
+    candidate = TreeToJax(TreeConventions()).apply(unit, facts, {"root": str(tmp_path)})
+    ported = candidate.files[Path("floor_mod_jax.py")].decode()
+    assert "floor = None" in ported
+    out = tmp_path / "emitted"
+    out.mkdir()
+    for path, content in candidate.files.items():
+        (out / path.name).write_bytes(content)
+    sys.path.insert(0, str(out))
+    try:
+        module = importlib.import_module("floor_mod_jax")
+        import numpy as np
+
+        coef = np.full((2, 3), 2.0, order="F")
+        x = np.ones((2, 3), order="F")
+        assert "apply_flat" in candidate.notes["jax"]["kernels"]
+        result = np.asarray(module.apply_flat(3, 2, x, 3, coef))
+        assert result.tolist() == [[2.0, 2.0, 2.0], [2.0, 2.0, 2.0]]
+    finally:
+        sys.path.remove(str(out))
+        for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"):
+            sys.modules.pop(f"floor_mod{suffix}", None)

From 614f0c5e72d130caafa61d1008c46ed069f082d2 Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 15:18:18 -0600
Subject: [PATCH 35/73] JAX lowering: a loop carries only names bound before
 it; subscripted tuple targets; elemental calls of kernels

* The anchor's _out = callee(...) result tuple inside a loop, unpacked
  on the next lines, was carried through the fori_loop and read for the
  initial carry before any assignment (CLUBB's new_hybrid_pdf_driver). A
  name not bound before the loop is the body's own.
* lo[i - 1, :], hi[i - 1, :] = split(...) goes through a temporary,
  each element a store of its own, instead of delegating the subprogram.
* _f_ecall(kernel, ...) is the runtime's jnp.vectorize over the
  kernel's implementation, its state closure appended; of a subprogram not
  being emitted it queues the caller, as a direct call does.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/jax/backend.py | 58 ++++++++++++++++++++++++++---
 src/recast/transform/jax/runtime.py |  7 ++++
 tests/test_jax_transform.py         | 53 ++++++++++++++++++++++++++
 3 files changed, 113 insertions(+), 5 deletions(-)

diff --git a/src/recast/transform/jax/backend.py b/src/recast/transform/jax/backend.py
index cdee415..ee251de 100644
--- a/src/recast/transform/jax/backend.py
+++ b/src/recast/transform/jax/backend.py
@@ -271,6 +271,22 @@ def __init__(self, call_map, known_subs):
     def visit_Call(self, node):
         self.generic_visit(node)
         f = node.func
+        if isinstance(f, ast.Name) and f.id == "_f_ecall" and node.args:
+            # ``_f_ecall(split, x[i - 1, :])``: the elemental broadcast of an
+            # emitted kernel is the runtime's jnp.vectorize over the kernel's
+            # implementation, its state closure appended; of a subprogram
+            # not being emitted, a host function no tracer can run.
+            callee = node.args[0]
+            if isinstance(callee, ast.Name) and callee.id in self.map:
+                info = self.map[callee.id]
+                node.args = [
+                    ast.Name(id=f"_{callee.id}_k_impl", ctx=ast.Load()),
+                    *node.args[1:],
+                    *[ast.Name(id=c, ctx=ast.Load()) for c in info["closure"]],
+                ]
+            elif isinstance(callee, ast.Name) and callee.id in self.subs:
+                raise JaxQueue(f"elemental call of non-emitted subprogram {callee.id}")
+            return node
         if isinstance(f, ast.Name):
             if f.id in self.map:
                 info = self.map[f.id]
@@ -481,8 +497,19 @@ class KernelLowerer:
         ast.ImportFrom,
     )
 
-    def __init__(self):
+    def __init__(self, bound=()):
         self.n = 0
+        # The names bound so far, in statement order: a loop carries only a
+        # name that exists before it. One first assigned inside the body
+        # (the anchor's ``_out = callee(...)`` result tuple, unpacked on the
+        # next lines) is the body's own -- carried, its initial value would
+        # be read before any assignment, and its shape may change between
+        # two calls in the same body.
+        self.bound = set(bound)
+
+    def _bind(self, stmts):
+        for name in _assigned_names(stmts):
+            self.bound.add(name)
 
     def lower_block(self, stmts, depth):
         out = []
@@ -502,9 +529,11 @@ def lower_block(self, stmts, depth):
             elif isinstance(s, ast.If):
                 out.extend(self.lower_if(s, depth))
             elif isinstance(s, ast.Assign):
-                out.append(self.lower_assign(s))
+                lowered = self.lower_assign(s)
+                out.extend(lowered if isinstance(lowered, list) else [lowered])
             else:
                 out.append(s)  # Expr (docstring), Return at top level, Pass
+            self._bind([s])
         return out
 
     def lower_assign(self, s):
@@ -515,7 +544,20 @@ def lower_assign(self, s):
             # multi-output intra-module call: a, b = _callee_k_impl(...)
             if all(isinstance(e, ast.Name) for e in t.elts):
                 return s
-            raise JaxQueue("tuple target with non-name elements")
+            # ``lo[i - 1, :], hi[i - 1, :] = split(...)`` (CLUBB's column
+            # loops): the tuple through a temporary, each element a store
+            # of its own -- a subscript store is lowered on its own line.
+            self.n += 1
+            tmp = f"_t{self.n}"
+            stores = [ast.Assign(targets=[ast.Name(id=tmp, ctx=ast.Store())], value=s.value)]
+            for at, element in enumerate(t.elts):
+                piece = ast.Subscript(
+                    value=ast.Name(id=tmp, ctx=ast.Load()),
+                    slice=ast.Constant(value=at),
+                    ctx=ast.Load(),
+                )
+                stores.append(self.lower_assign(ast.Assign(targets=[element], value=piece)))
+            return stores
         if isinstance(t, ast.Name):
             # strengthen scalar literal inits so fori_loop carries keep a
             # stable strong dtype (0.0 -> jnp.float64(0.0))
@@ -608,9 +650,15 @@ def lower_for(self, s, depth):
         else:
             raise JaxQueue("malformed range")
 
+        bound_before = set(self.bound)
+        self.bound.add(s.target.id)
         body = self.lower_block(_cycle_to_else(s.body, []), depth + 1)
         settled = _trace_constant_stores(body)
-        carried = [n for n in _assigned_names(body) if n != s.target.id and n not in settled]
+        carried = [
+            n
+            for n in _assigned_names(body)
+            if n != s.target.id and n not in settled and n in bound_before
+        ]
         if not carried:
             raise JaxQueue("loop with no carried effects")
         self.n += 1
@@ -797,7 +845,7 @@ def visit_FunctionDef(self, node):
     _bind_writer_calls(fn, call_map or {})
     ExprMap().visit(fn)
     CallRewrite(call_map or {}, known_subs or set()).visit(fn)
-    fn.body = KernelLowerer().lower_block(fn.body, 0)
+    fn.body = KernelLowerer(bound={a.arg for a in fn.args.args}).lower_block(fn.body, 0)
     ast.fix_missing_locations(fn)
     return ast.unparse(fn)
 
diff --git a/src/recast/transform/jax/runtime.py b/src/recast/transform/jax/runtime.py
index b162a65..6949629 100644
--- a/src/recast/transform/jax/runtime.py
+++ b/src/recast/transform/jax/runtime.py
@@ -38,6 +38,7 @@
 __all__ = [
     "_f_adjustl",
     "_f_dim",
+    "_f_ecall",
     "_f_epsilon",
     "_f_fori",
     "_f_huge",
@@ -165,6 +166,12 @@ def _f_fori(lo, hi, body, init):
     return lax.fori_loop(lo, hi, body, init)
 
 
+def _f_ecall(fn, *args, **kw):
+    """ELEMENTAL procedure broadcast over array actuals, as the NumPy
+    runtime does with np.vectorize: the scalar kernel per element."""
+    return jnp.vectorize(fn)(*args, **kw)
+
+
 def _f_sqrt(x):
     """Fortran SQRT: a NaN for a negative real, not an exception, and the
     correctly rounded root otherwise -- what the NumPy shim does with
diff --git a/tests/test_jax_transform.py b/tests/test_jax_transform.py
index 6dbe553..5ca5fd9 100644
--- a/tests/test_jax_transform.py
+++ b/tests/test_jax_transform.py
@@ -470,3 +470,56 @@ def test_a_flat_kernel_spells_the_optional_dummy_the_plan_leaves_out(tmp_path: P
         sys.path.remove(str(out))
         for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"):
             sys.modules.pop(f"floor_mod{suffix}", None)
+
+
+PAIRS_IN_A_LOOP = """\
+module pair_demo
+  implicit none
+  integer, parameter :: r8 = selected_real_kind(12)
+contains
+  elemental subroutine split(x, lo, hi)
+    real(r8), intent(in)  :: x
+    real(r8), intent(out) :: lo, hi
+    lo = x - 1.0_r8
+    hi = x + 1.0_r8
+  end subroutine split
+  subroutine bracket(n, nz, x, lo, hi)
+    integer,  intent(in)  :: n, nz
+    real(r8), intent(in)  :: x(n, nz)
+    real(r8), intent(out) :: lo(n, nz), hi(n, nz)
+    integer :: i
+    do i = 1, n
+      call split( x(i, :), lo(i, :), hi(i, :) )
+    end do
+  end subroutine bracket
+end module pair_demo
+"""
+
+
+def test_a_call_result_tuple_inside_a_loop_is_the_bodys_own(tmp_path: Path) -> None:
+    """The anchor spells a two-output call as ``_out = split(...)`` and
+    unpacks it on the next lines. Carried through the fori_loop, ``_out``
+    was read for the initial carry before any assignment
+    (UnboundLocalError; CLUBB's new_hybrid_pdf_driver). A name not bound
+    before the loop is the body's own."""
+    import importlib
+    import sys
+
+    candidate = port(tmp_path, PAIRS_IN_A_LOOP, "pair_demo")
+    assert "bracket" in candidate.notes["jax"]["kernels"]
+    out = tmp_path / "emitted"
+    out.mkdir()
+    for path, content in candidate.files.items():
+        (out / path.name).write_bytes(content)
+    sys.path.insert(0, str(out))
+    try:
+        module = importlib.import_module("pair_demo_jax")
+        import numpy as np
+
+        x = np.array([[1.0, 2.0], [3.0, 4.0]], order="F")
+        lo, hi = (np.asarray(v) for v in module.bracket(2, 2, x))
+        assert lo.tolist() == [[0.0, 1.0], [2.0, 3.0]] and hi.tolist() == [[2.0, 3.0], [4.0, 5.0]]
+    finally:
+        sys.path.remove(str(out))
+        for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"):
+            sys.modules.pop(f"pair_demo{suffix}", None)

From 2f2aec688540e31e7fefb3af16db315afeea809b Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 15:18:37 -0600
Subject: [PATCH 36/73] Rename the tuple-store index (mypy)

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/jax/backend.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/recast/transform/jax/backend.py b/src/recast/transform/jax/backend.py
index ee251de..5dac8f1 100644
--- a/src/recast/transform/jax/backend.py
+++ b/src/recast/transform/jax/backend.py
@@ -550,10 +550,10 @@ def lower_assign(self, s):
             self.n += 1
             tmp = f"_t{self.n}"
             stores = [ast.Assign(targets=[ast.Name(id=tmp, ctx=ast.Store())], value=s.value)]
-            for at, element in enumerate(t.elts):
+            for index, element in enumerate(t.elts):
                 piece = ast.Subscript(
                     value=ast.Name(id=tmp, ctx=ast.Load()),
-                    slice=ast.Constant(value=at),
+                    slice=ast.Constant(value=index),
                     ctx=ast.Load(),
                 )
                 stores.append(self.lower_assign(ast.Assign(targets=[element], value=piece)))

From 89f16a2b4f0db68c360061b4f5382c8abe97d0cd Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 15:20:10 -0600
Subject: [PATCH 37/73] JAX lowering: bound names come from what was emitted; a
 NumPy integer bound is static too

A loop's body-local (the anchor's _out) was bound for the next loop by the
original statement's assignments; the lowered loop binds its carried names
only. And a static kernel argument arrives as np.int32, which _f_fori did
not take for a static bound (CLUBB's sclr_dim = 0 tracer loops).

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/jax/backend.py | 6 +++++-
 src/recast/transform/jax/runtime.py | 4 +++-
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/src/recast/transform/jax/backend.py b/src/recast/transform/jax/backend.py
index 5dac8f1..67301b6 100644
--- a/src/recast/transform/jax/backend.py
+++ b/src/recast/transform/jax/backend.py
@@ -514,6 +514,7 @@ def _bind(self, stmts):
     def lower_block(self, stmts, depth):
         out = []
         for s in stmts:
+            start = len(out)
             if isinstance(s, self.BANNED):
                 raise JaxQueue(f"unsupported stmt {type(s).__name__}")
             if isinstance(s, ast.Continue | ast.Break):
@@ -533,7 +534,10 @@ def lower_block(self, stmts, depth):
                 out.extend(lowered if isinstance(lowered, list) else [lowered])
             else:
                 out.append(s)  # Expr (docstring), Return at top level, Pass
-            self._bind([s])
+            # Bound by what was *emitted*: a lowered loop or branch binds its
+            # carried names and nothing else -- a body-local of one loop is
+            # not bound for the next.
+            self._bind(out[start:])
         return out
 
     def lower_assign(self, s):
diff --git a/src/recast/transform/jax/runtime.py b/src/recast/transform/jax/runtime.py
index 6949629..714eac5 100644
--- a/src/recast/transform/jax/runtime.py
+++ b/src/recast/transform/jax/runtime.py
@@ -28,6 +28,7 @@
 """
 
 import jax
+import numpy as np
 
 jax.config.update("jax_enable_x64", True)
 
@@ -161,7 +162,8 @@ def _f_fori(lo, hi, body, init):
     trace the body once, and JAX refuses any index into a size-0 axis at
     trace time. A dynamic bound is left to ``fori_loop``.
     """
-    if isinstance(lo, int) and isinstance(hi, int) and hi <= lo:
+    static = (int, np.integer)
+    if isinstance(lo, static) and isinstance(hi, static) and int(hi) <= int(lo):
         return init
     return lax.fori_loop(lo, hi, body, init)
 

From 41d9f3474b82532b6ef6a213cf8fd2e97cd4c31b Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 15:20:24 -0600
Subject: [PATCH 38/73] Annotate the lowered statement list (mypy)

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/jax/backend.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/recast/transform/jax/backend.py b/src/recast/transform/jax/backend.py
index 67301b6..7e2398f 100644
--- a/src/recast/transform/jax/backend.py
+++ b/src/recast/transform/jax/backend.py
@@ -512,7 +512,7 @@ def _bind(self, stmts):
             self.bound.add(name)
 
     def lower_block(self, stmts, depth):
-        out = []
+        out: list[ast.stmt] = []
         for s in stmts:
             start = len(out)
             if isinstance(s, self.BANNED):

From f01cd80e55afd0d5f38ab6c80a393648fa2213fc Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 15:28:45 -0600
Subject: [PATCH 39/73] JAX port: a bare call statement binds the kernel's
 returned buffer outputs

CLUBB declares xp3_lg_2005_ansatz's xp3 without an intent and the
extension's frontend overrides it to OUT. The caller's anchor, seeing no
intent at the call, calls bare: the callee writes in place and returns
the buffer, the caller ignores the return. A kernel cannot write in place
-- its return is the output -- so the statement binds what comes back to
the actuals it was passed (advance_xp3 got back the zeros it passed in
for every xp3). The anchor's _out call-result tuple is never a carry.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/jax/backend.py |  7 ++-
 src/recast/transform/jax/tree.py    | 52 +++++++++++++++++++
 tests/test_jax_transform.py         | 79 +++++++++++++++++++++++++++++
 3 files changed, 136 insertions(+), 2 deletions(-)

diff --git a/src/recast/transform/jax/backend.py b/src/recast/transform/jax/backend.py
index 7e2398f..afdf8fb 100644
--- a/src/recast/transform/jax/backend.py
+++ b/src/recast/transform/jax/backend.py
@@ -361,11 +361,14 @@ def add(n):
                 # body before their use and never read after it; carried, they
                 # would have to be initialized at the enclosing level, where
                 # nothing assigns them.
-                if isinstance(t, ast.Name) and not re.fullmatch(r"_(?:hi_|cnt_|t)\d+", t.id):
+                # ``_out`` is the anchor's call-result tuple, assigned and
+                # unpacked on consecutive lines of one block: the block's own,
+                # never a carry (its shape changes from call to call).
+                if isinstance(t, ast.Name) and not re.fullmatch(r"_(?:hi_|cnt_|t)\d+|_out", t.id):
                     add(t.id)
                 elif isinstance(t, ast.Tuple):
                     for e in t.elts:
-                        if isinstance(e, ast.Name) and not re.fullmatch(r"_t\d+", e.id):
+                        if isinstance(e, ast.Name) and not re.fullmatch(r"_t\d+|_out", e.id):
                             add(e.id)
         elif isinstance(s, ast.If):
             for n in _assigned_names(s.body) + _assigned_names(s.orelse):
diff --git a/src/recast/transform/jax/tree.py b/src/recast/transform/jax/tree.py
index 14f6eca..bcef0f0 100644
--- a/src/recast/transform/jax/tree.py
+++ b/src/recast/transform/jax/tree.py
@@ -425,6 +425,10 @@ def visit_Expr(self, node: ast.Expr) -> Any:
         call = node.value
         if isinstance(call, ast.Call) and self._flat_callee(call) is not None:
             return self._rewrite_call(None, call)
+        if isinstance(call, ast.Call):
+            bound = self._bind_buffer_outputs(node, call)
+            if bound is not None:
+                return bound
         # ``_f_copy_out(dst, src)``: an in-place copy is an assignment here.
         if (
             isinstance(call, ast.Call)
@@ -753,6 +757,54 @@ def _companion_writes(self, node: ast.Call) -> list[str]:
             flats.append(flat)
         return flats
 
+    def _callee_record(self, call: ast.Call) -> dict[str, Any] | None:
+        """The interface record of a kernel a call statement reaches: one of
+        this module's, or a companion port's."""
+        func = call.func
+        if isinstance(func, ast.Name):
+            return (self.own.get("records") or {}).get(func.id)
+        if isinstance(func, ast.Attribute) and isinstance(func.value, ast.Name):
+            module = self.spelling.modules.get(func.value.id)
+            port = self.ports.get(module) if module is not None else None
+            if port is None or func.attr not in port.get("kernels", ()):
+                return None
+            return (port.get("records") or {}).get(func.attr)
+        return None
+
+    def _bind_buffer_outputs(self, node: ast.Expr, call: ast.Call) -> ast.stmt | None:
+        """``callee(..., up3)`` as a statement: the anchor's callee wrote the
+        OUT array in place and returned it; the caller ignored the return.
+        A kernel cannot write in place -- its return *is* the output -- so
+        the statement binds what comes back to the actuals it was passed
+        (CLUBB's advance_xp3 calling skx_module's xp3_lg_2005_ansatz:
+        without this every xp3 came back as the zeros it went in as)."""
+        record = self._callee_record(call)
+        if record is None or record.get("kind") == "function":
+            return None
+        outs = [
+            (at, a)
+            for at, a in enumerate(record["args"])
+            if a["intent"] in ("OUT", "INOUT") and not a.get("optional")
+        ]
+        if not outs:
+            return None
+        by_keyword = {k.arg: k.value for k in call.keywords if k.arg}
+        targets: list[ast.expr] = []
+        for at, a in outs:
+            actual = call.args[at] if at < len(call.args) else by_keyword.get(_py(a["name"]))
+            if actual is None or not isinstance(actual, ast.Name | ast.Subscript):
+                return None
+            target = copy.deepcopy(actual)
+            target.ctx = ast.Store()
+            targets.append(target)
+        target_node: ast.expr = (
+            targets[0] if len(targets) == 1 else ast.Tuple(elts=targets, ctx=ast.Store())
+        )
+        # As an assignment, through the assignment's own path: the callee's
+        # spelling, its state closure and its written state bind there.
+        assign = ast.Assign(targets=[target_node], value=call)
+        return self.visit(ast.copy_location(assign, node))
+
     # -- calls into companions ------------------------------------------------
 
     def _companion_call(self, node: ast.Call) -> ast.AST:
diff --git a/tests/test_jax_transform.py b/tests/test_jax_transform.py
index 5ca5fd9..2430173 100644
--- a/tests/test_jax_transform.py
+++ b/tests/test_jax_transform.py
@@ -523,3 +523,82 @@ def test_a_call_result_tuple_inside_a_loop_is_the_bodys_own(tmp_path: Path) -> N
         sys.path.remove(str(out))
         for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"):
             sys.modules.pop(f"pair_demo{suffix}", None)
+
+
+WRITER = """\
+module writer_mod
+  implicit none
+contains
+  subroutine fill(n, x, y)
+    integer,  intent(in)  :: n
+    real(8),  intent(in)  :: x(n)
+    real(8) :: y(n)
+    y = x + 1.0d0
+  end subroutine fill
+end module writer_mod
+"""
+
+CALLS_WRITER = """\
+module caller_mod
+  use writer_mod, only: fill
+  implicit none
+contains
+  subroutine run(n, x, y)
+    integer,  intent(in)  :: n
+    real(8),  intent(in)  :: x(n)
+    real(8),  intent(out) :: y(n)
+    call fill( n, x, y )
+    y = y * 2.0d0
+  end subroutine run
+end module caller_mod
+"""
+
+
+def test_a_bare_call_binds_the_kernels_returned_buffer(tmp_path: Path) -> None:
+    """CLUBB declares ``xp3_lg_2005_ansatz``'s ``xp3`` without an intent and
+    the extension's frontend overrides it to OUT. The caller's anchor, seeing
+    no intent at the call, calls bare: the callee writes ``y`` in place and
+    returns it, and the caller ignores the return. A kernel cannot write in
+    place -- its return *is* the output -- so the statement binds what
+    comes back to the actual; without it advance_xp3 got back the zeros it
+    passed in for every xp3."""
+    import importlib
+    import sys
+
+    from recast.registry import REGISTRY
+    from recast.transform.jax.tree import TreeToJax
+    from recast.transform.numpy.tree import TreeConventions
+
+    def overriding_frontend(**_config: object) -> FortranFrontend:
+        return FortranFrontend(
+            flatten=True, buffer_out_arrays="all", intent_overrides={"fill": {"y": "OUT"}}
+        )
+
+    REGISTRY.register("frontend", "fortran-y-out", overriding_frontend, replace=True)
+    (tmp_path / "writer_mod.f90").write_text(WRITER)
+    (tmp_path / "caller_mod.f90").write_text(CALLS_WRITER)
+    frontend = overriding_frontend()
+    unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:caller_mod")
+    facts = frontend.analyze(unit, tmp_path)
+    conventions = TreeConventions(frontend="fortran-y-out")
+    candidate = TreeToJax(conventions).apply(unit, facts, {"root": str(tmp_path)})
+    assert "run" in candidate.notes["jax"]["kernels"], candidate.notes["jax"]
+    ported = candidate.files[Path("caller_mod_jax.py")].decode()
+    assert "_writer_mod.fill(n, x, y)" in candidate.files[Path("caller_mod_numpy.py")].decode()
+    assert "y = _writer_mod_jax._fill_k_impl(n, x, y)" in ported
+    out = tmp_path / "emitted"
+    out.mkdir()
+    for path, content in candidate.files.items():
+        (out / path.name).write_bytes(content)
+    sys.path.insert(0, str(out))
+    try:
+        module = importlib.import_module("caller_mod_jax")
+        import numpy as np
+
+        got = module.run(2, np.array([1.0, 2.0]), np.zeros(2))
+        assert np.asarray(got).tolist() == [4.0, 6.0]
+    finally:
+        sys.path.remove(str(out))
+        for name in list(sys.modules):
+            if name.startswith(("caller_mod", "writer_mod")):
+                sys.modules.pop(name, None)

From f9ba9e5d85ffd3e77b6e8e25f2706e805280229c Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 15:29:01 -0600
Subject: [PATCH 40/73] Type the bound statement (mypy)

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/jax/tree.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/recast/transform/jax/tree.py b/src/recast/transform/jax/tree.py
index bcef0f0..584d5f8 100644
--- a/src/recast/transform/jax/tree.py
+++ b/src/recast/transform/jax/tree.py
@@ -803,7 +803,8 @@ def _bind_buffer_outputs(self, node: ast.Expr, call: ast.Call) -> ast.stmt | Non
         # As an assignment, through the assignment's own path: the callee's
         # spelling, its state closure and its written state bind there.
         assign = ast.Assign(targets=[target_node], value=call)
-        return self.visit(ast.copy_location(assign, node))
+        bound: ast.stmt | None = self.visit(ast.copy_location(assign, node))
+        return bound
 
     # -- calls into companions ------------------------------------------------
 

From d8837d62fadb37985ae4fbbedb02593e541d3f84 Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 15:35:26 -0600
Subject: [PATCH 41/73] JAX port: an elemental call of a companion procedure
 broadcasts its kernel

CLUBB's new_hybrid_pdf_driver calls new_hybrid_pdf's elemental
calculate_mixture_fraction over column slices; the anchor's
_f_ecall(_new.calculate_mixture_fraction, ...) kept the host attribute
and the vectorize traced a NumPy function. The companion's kernel
implementation goes under the vectorize; a procedure the companion's port
delegated makes the caller not flat.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/jax/tree.py | 33 +++++++++++++++
 tests/test_jax_transform.py      | 73 ++++++++++++++++++++++++++++++++
 2 files changed, 106 insertions(+)

diff --git a/src/recast/transform/jax/tree.py b/src/recast/transform/jax/tree.py
index 584d5f8..9072329 100644
--- a/src/recast/transform/jax/tree.py
+++ b/src/recast/transform/jax/tree.py
@@ -677,6 +677,39 @@ def visit_Call(self, node: ast.Call) -> ast.AST:
             raise NotFlat(
                 f"{ast.unparse(node.func)} takes the object and is called inside an expression"
             )
+        if (
+            isinstance(node.func, ast.Name)
+            and node.func.id == "_f_ecall"
+            and node.args
+            and isinstance(node.args[0], ast.Attribute)
+            and isinstance(node.args[0].value, ast.Name)
+        ):
+            # ``_f_ecall(_new.calculate_mixture_fraction, ...)``: the
+            # elemental broadcast of a companion's procedure. Its kernel's
+            # implementation goes under the vectorize; a procedure the
+            # companion's port delegated is a host function no tracer can
+            # run, so the caller is not flat either.
+            callee = node.args[0]
+            module = self.spelling.modules.get(callee.value.id)
+            port = self.ports.get(module) if module is not None else None
+            if port is not None:
+                if callee.attr not in port["kernels"]:
+                    raise NotFlat(
+                        f"elemental call of {module}.{callee.attr}, which its port did not lower"
+                    )
+                if (port.get("closures") or {}).get(callee.attr):
+                    raise NotFlat(
+                        f"elemental call of {module}.{callee.attr}, which reads module state"
+                    )
+                self.companions.add(callee.value.id)
+                node.args[0] = ast.copy_location(
+                    ast.Attribute(
+                        value=ast.Name(id=f"{callee.value.id}_jax", ctx=ast.Load()),
+                        attr=f"_{callee.attr}_k_impl",
+                        ctx=ast.Load(),
+                    ),
+                    callee,
+                )
         self.generic_visit(node)
         # ``int(x)`` and ``np.float64(x)`` on a traced value: the cast the
         # anchor spells with a Python or NumPy constructor is ``jnp``'s here.
diff --git a/tests/test_jax_transform.py b/tests/test_jax_transform.py
index 2430173..28dd9b9 100644
--- a/tests/test_jax_transform.py
+++ b/tests/test_jax_transform.py
@@ -602,3 +602,76 @@ def overriding_frontend(**_config: object) -> FortranFrontend:
         for name in list(sys.modules):
             if name.startswith(("caller_mod", "writer_mod")):
                 sys.modules.pop(name, None)
+
+
+ELEMENTAL_COMPANION = """\
+module elem_mod
+  implicit none
+  integer, parameter :: r8 = selected_real_kind(12)
+contains
+  elemental subroutine split(x, lo, hi)
+    real(r8), intent(in)  :: x
+    real(r8), intent(out) :: lo, hi
+    lo = x - 1.0_r8
+    hi = x + 1.0_r8
+  end subroutine split
+end module elem_mod
+"""
+
+CALLS_ELEMENTAL_COMPANION = """\
+module bracket_mod
+  use elem_mod, only: split
+  implicit none
+  integer, parameter :: r8 = selected_real_kind(12)
+contains
+  subroutine bracket(n, nz, x, lo, hi)
+    integer,  intent(in)  :: n, nz
+    real(r8), intent(in)  :: x(n, nz)
+    real(r8), intent(out) :: lo(n, nz), hi(n, nz)
+    integer :: i
+    do i = 1, n
+      call split( x(i, :), lo(i, :), hi(i, :) )
+    end do
+  end subroutine bracket
+end module bracket_mod
+"""
+
+
+def test_an_elemental_call_of_a_companion_broadcasts_its_kernel(tmp_path: Path) -> None:
+    """CLUBB's new_hybrid_pdf_driver calls new_hybrid_pdf's elemental
+    ``calculate_mixture_fraction`` over column slices: the anchor's
+    ``_f_ecall(_new.calculate_mixture_fraction, ...)``. Left as the host
+    attribute, the vectorize traced a NumPy function; the companion's kernel
+    implementation goes under it instead."""
+    import importlib
+    import sys
+
+    from recast.transform.jax.tree import TreeToJax
+    from recast.transform.numpy.tree import TreeConventions
+
+    (tmp_path / "elem_mod.f90").write_text(ELEMENTAL_COMPANION)
+    (tmp_path / "bracket_mod.f90").write_text(CALLS_ELEMENTAL_COMPANION)
+    frontend = FortranFrontend(flatten=True)
+    unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:bracket_mod")
+    facts = frontend.analyze(unit, tmp_path)
+    candidate = TreeToJax(TreeConventions()).apply(unit, facts, {"root": str(tmp_path)})
+    assert "bracket" in candidate.notes["jax"]["kernels"], candidate.notes["jax"]
+    ported = candidate.files[Path("bracket_mod_jax.py")].decode()
+    assert "_f_ecall(_elem_mod_jax._split_k_impl, " in ported
+    out = tmp_path / "emitted"
+    out.mkdir()
+    for path, content in candidate.files.items():
+        (out / path.name).write_bytes(content)
+    sys.path.insert(0, str(out))
+    try:
+        module = importlib.import_module("bracket_mod_jax")
+        import numpy as np
+
+        x = np.array([[1.0, 2.0], [3.0, 4.0]], order="F")
+        lo, hi = (np.asarray(v) for v in module.bracket(2, 2, x))
+        assert lo.tolist() == [[0.0, 1.0], [2.0, 3.0]] and hi.tolist() == [[2.0, 3.0], [4.0, 5.0]]
+    finally:
+        sys.path.remove(str(out))
+        for name in list(sys.modules):
+            if name.startswith(("bracket_mod", "elem_mod")):
+                sys.modules.pop(name, None)

From 1c84d38fcf49721c1b0c9328a5588823cfa6b64e Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 15:35:45 -0600
Subject: [PATCH 42/73] Type the elemental callee's owner (mypy)

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/jax/tree.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/src/recast/transform/jax/tree.py b/src/recast/transform/jax/tree.py
index 9072329..e5341d9 100644
--- a/src/recast/transform/jax/tree.py
+++ b/src/recast/transform/jax/tree.py
@@ -690,7 +690,9 @@ def visit_Call(self, node: ast.Call) -> ast.AST:
             # companion's port delegated is a host function no tracer can
             # run, so the caller is not flat either.
             callee = node.args[0]
-            module = self.spelling.modules.get(callee.value.id)
+            owner = callee.value
+            assert isinstance(owner, ast.Name)
+            module = self.spelling.modules.get(owner.id)
             port = self.ports.get(module) if module is not None else None
             if port is not None:
                 if callee.attr not in port["kernels"]:
@@ -701,10 +703,10 @@ def visit_Call(self, node: ast.Call) -> ast.AST:
                     raise NotFlat(
                         f"elemental call of {module}.{callee.attr}, which reads module state"
                     )
-                self.companions.add(callee.value.id)
+                self.companions.add(owner.id)
                 node.args[0] = ast.copy_location(
                     ast.Attribute(
-                        value=ast.Name(id=f"{callee.value.id}_jax", ctx=ast.Load()),
+                        value=ast.Name(id=f"{owner.id}_jax", ctx=ast.Load()),
                         attr=f"_{callee.attr}_k_impl",
                         ctx=ast.Load(),
                     ),

From 347f05bd78585afc702ac0becd06fce59811fc1a Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 15:41:18 -0600
Subject: [PATCH 43/73] JAX lowering: a branch on a static scalar argument is a
 trace-time if

CLUBB guards its scalar-tracer stores with if ( sclr_dim > 0 ).
Lowered to lax.cond both arms are traced, and the store into an array the
run never allocated is an IndexError at trace time. The kernel's static
scalar arguments are Python ints under jit, so a comparison over them
(and module constants) stays a Python if, as PRESENT tests already did.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/jax/backend.py | 36 ++++++++++++++++------
 tests/test_jax_transform.py         | 48 +++++++++++++++++++++++++++++
 2 files changed, 75 insertions(+), 9 deletions(-)

diff --git a/src/recast/transform/jax/backend.py b/src/recast/transform/jax/backend.py
index afdf8fb..627c708 100644
--- a/src/recast/transform/jax/backend.py
+++ b/src/recast/transform/jax/backend.py
@@ -376,10 +376,15 @@ def add(n):
     return out
 
 
-def _static_test(test):
+def _static_test(test, statics=frozenset()):
     """True for branch conditions decidable at trace time per the
     translate.py grammar: `x is [not] None` (Fortran PRESENT) and bare
-    `want_*` sentinels (optional-output flags, static under jit)."""
+    `want_*` sentinels (optional-output flags, static under jit), and a
+    comparison over module constants and the kernel's static scalar
+    arguments (``statics``: Python ints at trace time under jit). CLUBB's
+    ``if ( sclr_dim > 0 )`` guards stores into arrays the run never
+    allocated; lowered to lax.cond both arms are traced, and the store
+    into a (0, 0) array is an IndexError at trace time."""
     if (
         isinstance(test, ast.Compare)
         and len(test.ops) == 1
@@ -393,7 +398,7 @@ def constant(node):
         if isinstance(node, ast.Constant):
             return isinstance(node.value, (int, float)) and not isinstance(node.value, bool)
         if isinstance(node, ast.Name):
-            return node.id.isupper() and len(node.id) > 1
+            return node.id in statics or (node.id.isupper() and len(node.id) > 1)
         if isinstance(node, ast.BinOp):
             return constant(node.left) and constant(node.right)
         if isinstance(node, ast.UnaryOp):
@@ -406,7 +411,7 @@ def constant(node):
     if isinstance(test, ast.Compare) and len(test.ops) == 1:
         return constant(test.left) and all(constant(c) for c in test.comparators)
     if isinstance(test, ast.BoolOp):
-        return all(_static_test(v) for v in test.values)
+        return all(_static_test(v, statics) for v in test.values)
     return False
 
 
@@ -500,8 +505,9 @@ class KernelLowerer:
         ast.ImportFrom,
     )
 
-    def __init__(self, bound=()):
+    def __init__(self, bound=(), statics=()):
         self.n = 0
+        self.statics = frozenset(statics)
         # The names bound so far, in statement order: a loop carries only a
         # name that exists before it. One first assigned inside the body
         # (the anchor's ``_out = callee(...)`` result tuple, unpacked on the
@@ -726,7 +732,7 @@ def lower_if(self, s, depth):
         a None arg — is never traced)."""
         body = self.lower_block(s.body, depth + 1)
         orelse = self.lower_block(s.orelse, depth + 1)
-        if _static_test(s.test):
+        if _static_test(s.test, self.statics):
             return [ast.If(test=s.test, body=body, orelse=orelse or [])]
         carried = _assigned_names(body)
         for n in _assigned_names(orelse):
@@ -800,7 +806,9 @@ def split_params(fn_src):
     return params[:n_req], params[n_req:], fn_src.args.defaults
 
 
-def emit_kernel(fn_src, sub, closure, call_map=None, known_subs=None, writes=()):
+def emit_kernel(
+    fn_src, sub, closure, call_map=None, known_subs=None, writes=(), traced_scalars=frozenset()
+):
     """Original numpy FunctionDef -> unparsed _<name>_k_impl source.
 
     Kernel signature: [required..., closure..., optional-with-defaults].
@@ -852,7 +860,11 @@ def visit_FunctionDef(self, node):
     _bind_writer_calls(fn, call_map or {})
     ExprMap().visit(fn)
     CallRewrite(call_map or {}, known_subs or set()).visit(fn)
-    fn.body = KernelLowerer(bound={a.arg for a in fn.args.args}).lower_block(fn.body, 0)
+    nums, _ = static_spec(fn_src, sub, traced_scalars)
+    required, _, _ = split_params(fn_src)
+    statics = {required[pos] for pos in nums}
+    lowerer = KernelLowerer(bound={a.arg for a in fn.args.args}, statics=statics)
+    fn.body = lowerer.lower_block(fn.body, 0)
     ast.fix_missing_locations(fn)
     return ast.unparse(fn)
 
@@ -988,7 +1000,13 @@ def build_module(
             call_map = {k: v for k, v in call_map_all.items() if k != name}
             try:
                 srcs[name] = emit_kernel(
-                    fns[name], subs[name], closures[name], call_map, set(subs), wclosures[name]
+                    fns[name],
+                    subs[name],
+                    closures[name],
+                    call_map,
+                    set(subs),
+                    wclosures[name],
+                    traced_scalars,
                 )
             except JaxQueue as e:
                 failed[name] = f"[emit] {e}"
diff --git a/tests/test_jax_transform.py b/tests/test_jax_transform.py
index 28dd9b9..e06ec98 100644
--- a/tests/test_jax_transform.py
+++ b/tests/test_jax_transform.py
@@ -675,3 +675,51 @@ def test_an_elemental_call_of_a_companion_broadcasts_its_kernel(tmp_path: Path)
         for name in list(sys.modules):
             if name.startswith(("bracket_mod", "elem_mod")):
                 sys.modules.pop(name, None)
+
+
+GUARDED_BY_A_STATIC = """\
+module guard_demo
+  implicit none
+contains
+  subroutine first_tracer(n, m, x, y)
+    integer,  intent(in)  :: n, m
+    real(8),  intent(in)  :: x(n, m)
+    real(8),  intent(out) :: y(n)
+    y = 0.0d0
+    if ( m > 0 ) then
+      y(:) = x(:, 1)
+    end if
+  end subroutine first_tracer
+end module guard_demo
+"""
+
+
+def test_a_branch_on_a_static_scalar_is_a_trace_time_if(tmp_path: Path) -> None:
+    """CLUBB guards its scalar-tracer stores with ``if ( sclr_dim > 0 )``.
+    Lowered to lax.cond both arms are traced, and the store into the
+    zero-extent (or never allocated) array is an IndexError at trace time.
+    ``sclr_dim`` is a static argument of the kernel -- a Python int under
+    jit -- so the branch is a Python if."""
+    import importlib
+    import sys
+
+    candidate = port(tmp_path, GUARDED_BY_A_STATIC, "guard_demo")
+    assert candidate.notes["jax"]["kernels"] == ["first_tracer"]
+    ported = candidate.files[Path("guard_demo_jax.py")].decode()
+    assert "if m > 0:" in ported
+    out = tmp_path / "emitted"
+    out.mkdir()
+    for path, content in candidate.files.items():
+        (out / path.name).write_bytes(content)
+    sys.path.insert(0, str(out))
+    try:
+        module = importlib.import_module("guard_demo_jax")
+        import numpy as np
+
+        assert np.asarray(module.first_tracer(2, 0, np.zeros((2, 0)))).tolist() == [0.0, 0.0]
+        x = np.array([[1.0, 2.0], [3.0, 4.0]], order="F")
+        assert np.asarray(module.first_tracer(2, 2, x)).tolist() == [1.0, 3.0]
+    finally:
+        sys.path.remove(str(out))
+        for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"):
+            sys.modules.pop(f"guard_demo{suffix}", None)

From 19624d48f53659d4215be036d066ba28011a370b Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 15:49:20 -0600
Subject: [PATCH 44/73] JAX lowering: a branch on a static scalar is a Python
 if when concrete, a lax.cond when traced

A kernel's static scalar argument is a Python int through its jit wrapper
and a tracer when another kernel's traced body calls the implementation
(CLUBB's mono_cubic_interp takes its level indices as arguments, from its
caller's loop index). The emitted branch asks the runtime's _f_concrete at
trace time and takes the form that fits.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/jax/backend.py | 32 ++++++++++++++++++++++++++++-
 src/recast/transform/jax/runtime.py | 10 +++++++++
 tests/test_jax_transform.py         |  2 +-
 3 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/src/recast/transform/jax/backend.py b/src/recast/transform/jax/backend.py
index 627c708..ac59d94 100644
--- a/src/recast/transform/jax/backend.py
+++ b/src/recast/transform/jax/backend.py
@@ -732,8 +732,38 @@ def lower_if(self, s, depth):
         a None arg — is never traced)."""
         body = self.lower_block(s.body, depth + 1)
         orelse = self.lower_block(s.orelse, depth + 1)
-        if _static_test(s.test, self.statics):
+        if _static_test(s.test):
             return [ast.If(test=s.test, body=body, orelse=orelse or [])]
+        if _static_test(s.test, self.statics):
+            # Over the kernel's static scalar arguments: a Python if when
+            # the kernel is called through its jit wrapper (the arguments
+            # are Python ints), and the lax.cond when another kernel calls
+            # the implementation from a traced body (CLUBB's mono_cubic_interp
+            # takes its level indices as arguments, and its caller's loop
+            # index reaches it as a tracer). Decided at trace time by the
+            # runtime's ``_f_concrete``.
+            python_form = [ast.If(test=s.test, body=body, orelse=orelse or [])]
+            try:
+                cond_form = self._cond_form(s, copy.deepcopy(body), copy.deepcopy(orelse))
+            except JaxQueue:
+                return python_form  # nothing a cond could carry: the branch is a guard alone
+            if not cond_form:
+                return python_form
+            return [
+                ast.If(
+                    test=ast.Call(
+                        func=ast.Name(id="_f_concrete", ctx=ast.Load()),
+                        args=[copy.deepcopy(s.test)],
+                        keywords=[],
+                    ),
+                    body=python_form,
+                    orelse=cond_form,
+                )
+            ]
+        return self._cond_form(s, body, orelse)
+
+    def _cond_form(self, s, body, orelse):
+        """The lax.cond lowering of an if whose test is traced."""
         carried = _assigned_names(body)
         for n in _assigned_names(orelse):
             if n not in carried:
diff --git a/src/recast/transform/jax/runtime.py b/src/recast/transform/jax/runtime.py
index 714eac5..bc58786 100644
--- a/src/recast/transform/jax/runtime.py
+++ b/src/recast/transform/jax/runtime.py
@@ -38,6 +38,7 @@
 # A star-import from the generated module must see the underscore names.
 __all__ = [
     "_f_adjustl",
+    "_f_concrete",
     "_f_dim",
     "_f_ecall",
     "_f_epsilon",
@@ -154,6 +155,15 @@ def _f_vpow(a, b):
     return jnp.asarray(a) ** b
 
 
+def _f_concrete(x):
+    """Whether a value is known at trace time -- a Python or NumPy scalar,
+    or a concrete array -- rather than a tracer. A branch over a kernel's
+    static scalar argument is a Python ``if`` when the kernel runs through
+    its jit wrapper and a ``lax.cond`` when another kernel's traced body
+    calls its implementation; this decides which."""
+    return not isinstance(x, jax.core.Tracer)
+
+
 def _f_fori(lo, hi, body, init):
     """``lax.fori_loop`` unless the trip count is static and empty.
 
diff --git a/tests/test_jax_transform.py b/tests/test_jax_transform.py
index e06ec98..5cff841 100644
--- a/tests/test_jax_transform.py
+++ b/tests/test_jax_transform.py
@@ -706,7 +706,7 @@ def test_a_branch_on_a_static_scalar_is_a_trace_time_if(tmp_path: Path) -> None:
     candidate = port(tmp_path, GUARDED_BY_A_STATIC, "guard_demo")
     assert candidate.notes["jax"]["kernels"] == ["first_tracer"]
     ported = candidate.files[Path("guard_demo_jax.py")].decode()
-    assert "if m > 0:" in ported
+    assert "if _f_concrete(m > 0):" in ported and "if m > 0:" in ported
     out = tmp_path / "emitted"
     out.mkdir()
     for path, content in candidate.files.items():

From 781dde2cbccdbdf6ecf1101565eb4bf7e77d74fd Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 15:49:39 -0600
Subject: [PATCH 45/73] Type the Python form of the branch (mypy)

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/jax/backend.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/recast/transform/jax/backend.py b/src/recast/transform/jax/backend.py
index ac59d94..f598065 100644
--- a/src/recast/transform/jax/backend.py
+++ b/src/recast/transform/jax/backend.py
@@ -742,7 +742,7 @@ def lower_if(self, s, depth):
             # takes its level indices as arguments, and its caller's loop
             # index reaches it as a tracer). Decided at trace time by the
             # runtime's ``_f_concrete``.
-            python_form = [ast.If(test=s.test, body=body, orelse=orelse or [])]
+            python_form: list[ast.stmt] = [ast.If(test=s.test, body=body, orelse=orelse or [])]
             try:
                 cond_form = self._cond_form(s, copy.deepcopy(body), copy.deepcopy(orelse))
             except JaxQueue:

From e772aa404b7a9dd1a86ae3c41359e4204d29b68c Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 15:52:06 -0600
Subject: [PATCH 46/73] JAX runtime: an elemental call maps the kernel over the
 elements in sequence, not under vmap

jnp.vectorize under jit is a vmap: the kernel's branches become selects
and its loops batched, and CLUBB's hybrid PDF closure came out 1e7 ULP
from the NumPy anchor that way -- 2 ULP with lax.map, which keeps each
element's own control flow, as np.vectorize does on the anchor side.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/jax/runtime.py | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/src/recast/transform/jax/runtime.py b/src/recast/transform/jax/runtime.py
index bc58786..c28ee95 100644
--- a/src/recast/transform/jax/runtime.py
+++ b/src/recast/transform/jax/runtime.py
@@ -179,9 +179,18 @@ def _f_fori(lo, hi, body, init):
 
 
 def _f_ecall(fn, *args, **kw):
-    """ELEMENTAL procedure broadcast over array actuals, as the NumPy
-    runtime does with np.vectorize: the scalar kernel per element."""
-    return jnp.vectorize(fn)(*args, **kw)
+    """ELEMENTAL procedure broadcast over array actuals: the scalar kernel
+    per element, in sequence, as the NumPy runtime's np.vectorize runs it.
+
+    Not jnp.vectorize: under jit that is a vmap, which turns the kernel's
+    branches into selects and its loops into batched ones, and CLUBB's
+    hybrid PDF closure came out 1e7 ULP from the anchor that way (2 ULP
+    this way). lax.map keeps each element's own control flow."""
+    arrays = [jnp.asarray(a) for a in args]
+    shape = jnp.broadcast_shapes(*[a.shape for a in arrays])
+    flat = tuple(jnp.broadcast_to(a, shape).reshape(-1) for a in arrays)
+    outs = lax.map(lambda xs: fn(*xs, **kw), flat)
+    return jax.tree_util.tree_map(lambda o: o.reshape(shape), outs)
 
 
 def _f_sqrt(x):

From 756e0744eb6552803d52294026854b8e1aca6753 Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 16:43:32 -0600
Subject: [PATCH 47/73] JAX port: early returns of the same tuple fold into the
 branch structure

CLUBB's advance_clubb_core returns after each solver when the error code
says so, with the outputs as they stand. The single-exit rewrite merges a
single value with a where and refused a tuple, so the whole step was
host-delegated. When every early return is the final tuple, what follows
a returning branch moves into the branches that do not return, and the
kernel keeps one exit.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/jax/tree.py | 61 ++++++++++++++++++++++---
 tests/test_jax_transform.py      | 77 ++++++++++++++++++++++++++++++++
 2 files changed, 131 insertions(+), 7 deletions(-)

diff --git a/src/recast/transform/jax/tree.py b/src/recast/transform/jax/tree.py
index e5341d9..f8b8177 100644
--- a/src/recast/transform/jax/tree.py
+++ b/src/recast/transform/jax/tree.py
@@ -1584,6 +1584,39 @@ def _has_return(stmts: list[ast.stmt]) -> bool:
     return any(isinstance(n, ast.Return) for s in stmts for n in ast.walk(s))
 
 
+def _fold_returns(stmts: list[ast.stmt], rest: list[ast.stmt]) -> list[ast.stmt] | None:
+    """``stmts`` followed by ``rest`` with every early ``return`` folded into
+    the branch structure: the continuation moves into the branches that do
+    not return, so the function's one return at the end is reached by all
+    paths with the outputs as the early return would have left them. None
+    when a return sits inside a loop, which this fold cannot express."""
+    out: list[ast.stmt] = []
+    for at, statement in enumerate(stmts):
+        if isinstance(statement, ast.Return):
+            return out  # what follows is never reached on this path
+        if isinstance(statement, ast.For | ast.While) and _has_return([statement]):
+            return None
+        if isinstance(statement, ast.If) and _has_return([statement]):
+            # The continuation, itself folded: a return further down the
+            # same block would otherwise ride into the branch unfolded.
+            tail = _fold_returns(stmts[at + 1 :], rest)
+            if tail is None:
+                return None
+            body = _fold_returns(statement.body, copy.deepcopy(tail))
+            orelse = _fold_returns(statement.orelse, copy.deepcopy(tail))
+            if body is None or orelse is None:
+                return None
+            return [
+                *out,
+                ast.copy_location(
+                    ast.If(test=statement.test, body=body or [ast.Pass()], orelse=orelse),
+                    statement,
+                ),
+            ]
+        out.append(statement)
+    return [*out, *rest]
+
+
 def _single_exit(body: list[ast.stmt]) -> list[ast.stmt]:
     """Early returns -- ``if f0 == 0: root = x0; return root`` -- become a
     flag and a value, every later statement runs under ``if not _ret``, and
@@ -1593,13 +1626,27 @@ def _single_exit(body: list[ast.stmt]) -> list[ast.stmt]:
     if not early or not isinstance(body[-1], ast.Return) or body[-1].value is None:
         return body
     # The merge is one ``jnp.where`` over one value: a function with several
-    # outputs returns a tuple, which ``where`` cannot select.
-    for statement in body:
-        for node in ast.walk(statement):
-            if isinstance(node, ast.Return) and isinstance(node.value, ast.Tuple):
-                raise NotFlat(
-                    f"an early return in a function with several outputs: {ast.unparse(node)}"
-                )
+    # outputs returns a tuple, which ``where`` cannot select. When every
+    # early return hands back the same tuple the final one does (CLUBB's
+    # advance_clubb_core: ``if ( fatal ) return`` after each solver, the
+    # outputs as they stand), the returns fold into the branch structure
+    # instead: what follows a returning branch moves into the branches that
+    # do not return, and the one return at the end is reached by all.
+    tuples = [
+        node
+        for statement in body
+        for node in ast.walk(statement)
+        if isinstance(node, ast.Return) and isinstance(node.value, ast.Tuple)
+    ]
+    if tuples:
+        final = body[-1]
+        same = all(ast.dump(node.value) == ast.dump(final.value) for node in tuples)
+        folded = _fold_returns(body[:-1], []) if same else None
+        if folded is None:
+            raise NotFlat(
+                f"an early return in a function with several outputs: {ast.unparse(tuples[0])}"
+            )
+        return [*folded, final]
 
     class Returns(ast.NodeTransformer):
         def visit_Return(self, node: ast.Return) -> Any:
diff --git a/tests/test_jax_transform.py b/tests/test_jax_transform.py
index 5cff841..a251fac 100644
--- a/tests/test_jax_transform.py
+++ b/tests/test_jax_transform.py
@@ -723,3 +723,80 @@ def test_a_branch_on_a_static_scalar_is_a_trace_time_if(tmp_path: Path) -> None:
         sys.path.remove(str(out))
         for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"):
             sys.modules.pop(f"guard_demo{suffix}", None)
+
+
+EARLY_RETURNS = """\
+module early_mod
+  implicit none
+  type coefs_type
+    real(8), allocatable :: coef(:, :)
+  end type coefs_type
+contains
+  subroutine init_coefs( nz, ngrdcol, c )
+    integer, intent(in) :: nz, ngrdcol
+    type(coefs_type), intent(inout) :: c
+    allocate( c%coef(1:ngrdcol, 1:nz) )
+    c%coef = 2.0d0
+  end subroutine init_coefs
+  subroutine apply( nzt, ngrdcol, c, x, y, bad, worse )
+    integer, intent(in) :: nzt, ngrdcol
+    type(coefs_type), intent(in) :: c
+    real(8), intent(inout), dimension(ngrdcol, nzt) :: x
+    real(8), intent(out), dimension(ngrdcol, nzt) :: y
+    logical, intent(in) :: bad, worse
+    x = x * c%coef(:, 1:nzt)
+    y = x
+    if ( bad ) then
+      if ( worse ) then
+        return
+      end if
+      y = y + 1.0d0
+      return
+    end if
+    y = y + 10.0d0
+  end subroutine apply
+end module early_mod
+"""
+
+
+def test_early_returns_of_the_same_tuple_fold_into_the_branches(tmp_path: Path) -> None:
+    """CLUBB's advance_clubb_core returns after each solver when the error
+    code says so, the outputs as they stand. The single-exit rewrite merges
+    one value with a where and refused a tuple; when every early return is
+    the final tuple, the continuation folds into the non-returning branches
+    and the kernel keeps one exit."""
+    import importlib
+    import sys
+
+    from recast.transform.jax.tree import TreeToJax
+    from recast.transform.numpy.tree import TreeConventions
+
+    (tmp_path / "early_mod.f90").write_text(EARLY_RETURNS)
+    frontend = FortranFrontend(flatten={"patch_count": "ngrdcol", "bounds_pattern": r"^ngrdcol$"})
+    unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:early_mod")
+    facts = frontend.analyze(unit, tmp_path)
+    candidate = TreeToJax(TreeConventions()).apply(unit, facts, {"root": str(tmp_path)})
+    assert "apply_flat" in candidate.notes["jax"]["kernels"], candidate.notes["jax"]
+    out = tmp_path / "emitted"
+    out.mkdir()
+    for path, content in candidate.files.items():
+        (out / path.name).write_bytes(content)
+    sys.path.insert(0, str(out))
+    try:
+        module = importlib.import_module("early_mod_jax")
+        import numpy as np
+
+        coef = np.full((1, 2), 2.0, order="F")
+
+        def run(bad, worse):
+            x = np.ones((1, 2), order="F")
+            _x, y = module.apply_flat(2, 1, x, np.zeros((1, 2), order="F"), bad, worse, 2, coef)
+            return np.asarray(y).tolist()
+
+        assert run(False, False) == [[12.0, 12.0]]
+        assert run(True, False) == [[3.0, 3.0]]
+        assert run(True, True) == [[2.0, 2.0]]
+    finally:
+        sys.path.remove(str(out))
+        for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"):
+            sys.modules.pop(f"early_mod{suffix}", None)

From 86055fa657dbee311b8c539adecbd5a0ed4b6b06 Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 16:45:56 -0600
Subject: [PATCH 48/73] JAX port: an if whose body is empty keeps its else
 under the inverted guard

The fold of an early return leaves if returned: pass else: <the rest>;
the simplification for log-only branches ran the rest unconditionally.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/jax/tree.py | 20 +++++++++++++++++---
 1 file changed, 17 insertions(+), 3 deletions(-)

diff --git a/src/recast/transform/jax/tree.py b/src/recast/transform/jax/tree.py
index f8b8177..a8e4184 100644
--- a/src/recast/transform/jax/tree.py
+++ b/src/recast/transform/jax/tree.py
@@ -644,8 +644,20 @@ def visit_If(self, node: ast.If) -> Any:
         self.generic_visit(node)
         if all(isinstance(s, (ast.Pass, ast.Expr)) for s in node.body):
             # ``if cond: write(iulog, ...)`` -- a log line the anchor already
-            # left as ``pass``; nothing to carry.
-            return [*node.orelse] if node.orelse else None
+            # left as ``pass``; nothing to carry. An else branch keeps its
+            # guard, inverted: the fold of an early return leaves exactly
+            # ``if returned: pass else: <the rest>``, and running the rest
+            # unconditionally was the returned path's outputs overwritten.
+            if not node.orelse:
+                return None
+            return ast.copy_location(
+                ast.If(
+                    test=ast.UnaryOp(op=ast.Not(), operand=node.test),
+                    body=list(node.orelse),
+                    orelse=[],
+                ),
+                node,
+            )
         return node
 
     def visit_Raise(self, node: ast.Raise) -> Any:
@@ -1640,7 +1652,9 @@ def _single_exit(body: list[ast.stmt]) -> list[ast.stmt]:
     ]
     if tuples:
         final = body[-1]
-        same = all(ast.dump(node.value) == ast.dump(final.value) for node in tuples)
+        same = final.value is not None and all(
+            ast.dump(node.value) == ast.dump(final.value) for node in tuples
+        )
         folded = _fold_returns(body[:-1], []) if same else None
         if folded is None:
             raise NotFlat(

From 6b1b063d5c831204cc0fb8fcb9803e9a03ef4b70 Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 16:46:18 -0600
Subject: [PATCH 49/73] Narrow the return values before comparing them (mypy)

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/jax/tree.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/recast/transform/jax/tree.py b/src/recast/transform/jax/tree.py
index a8e4184..070ed47 100644
--- a/src/recast/transform/jax/tree.py
+++ b/src/recast/transform/jax/tree.py
@@ -1652,8 +1652,10 @@ def _single_exit(body: list[ast.stmt]) -> list[ast.stmt]:
     ]
     if tuples:
         final = body[-1]
-        same = final.value is not None and all(
-            ast.dump(node.value) == ast.dump(final.value) for node in tuples
+        final_value = final.value
+        same = final_value is not None and all(
+            node.value is not None and ast.dump(node.value) == ast.dump(final_value)
+            for node in tuples
         )
         folded = _fold_returns(body[:-1], []) if same else None
         if folded is None:

From 62a97abc7252c57ce74e92bae1f0deee9ee17ed3 Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 17:37:07 -0600
Subject: [PATCH 50/73] JAX port: an early return inside a loop becomes a flag

advance_clubb_core checks the error code per column inside a loop and
returns. No branch structure holds that: the return sets a traced flag,
every later statement at every level runs under it, the loop's remaining
iterations do nothing, and the one return at the end hands back the
outputs as the early return left them.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/jax/tree.py | 53 ++++++++++++++++++++--
 tests/test_jax_transform.py      | 77 ++++++++++++++++++++++++++++++++
 2 files changed, 127 insertions(+), 3 deletions(-)

diff --git a/src/recast/transform/jax/tree.py b/src/recast/transform/jax/tree.py
index 070ed47..1dfd92a 100644
--- a/src/recast/transform/jax/tree.py
+++ b/src/recast/transform/jax/tree.py
@@ -1629,6 +1629,39 @@ def _fold_returns(stmts: list[ast.stmt], rest: list[ast.stmt]) -> list[ast.stmt]
     return [*out, *rest]
 
 
+def _guard_after_returns(stmts: list[ast.stmt]) -> list[ast.stmt]:
+    """``return`` -> ``_ret = True``; what follows a statement that may have
+    returned, in the same block, runs under ``if not _ret`` -- recursively
+    through branches and loop bodies."""
+    out: list[ast.stmt] = []
+    exited = False
+    for statement in stmts:
+        if isinstance(statement, ast.Return):
+            rewritten: ast.stmt = ast.copy_location(
+                ast.Assign(
+                    targets=[ast.Name(id="_ret", ctx=ast.Store())],
+                    value=_jnp("bool_", [ast.Constant(True)]),
+                ),
+                statement,
+            )
+        else:
+            rewritten = statement
+            for field in ("body", "orelse"):
+                inner = getattr(statement, field, None)
+                if isinstance(inner, list) and inner:
+                    setattr(statement, field, _guard_after_returns(inner))
+        if exited:
+            rewritten = ast.If(
+                test=ast.UnaryOp(op=ast.Not(), operand=ast.Name(id="_ret", ctx=ast.Load())),
+                body=[rewritten],
+                orelse=[],
+            )
+        out.append(rewritten)
+        if _has_return([statement]):
+            exited = True
+    return out
+
+
 def _single_exit(body: list[ast.stmt]) -> list[ast.stmt]:
     """Early returns -- ``if f0 == 0: root = x0; return root`` -- become a
     flag and a value, every later statement runs under ``if not _ret``, and
@@ -1657,12 +1690,26 @@ def _single_exit(body: list[ast.stmt]) -> list[ast.stmt]:
             node.value is not None and ast.dump(node.value) == ast.dump(final_value)
             for node in tuples
         )
-        folded = _fold_returns(body[:-1], []) if same else None
-        if folded is None:
+        if not same:
             raise NotFlat(
                 f"an early return in a function with several outputs: {ast.unparse(tuples[0])}"
             )
-        return [*folded, final]
+        folded = _fold_returns(body[:-1], [])
+        if folded is not None:
+            return [*folded, final]
+        # A return inside a loop (advance_clubb_core checks the error code
+        # per column): a flag instead. The return sets it, every later
+        # statement at every level runs under ``if not _ret``, and the
+        # loop's remaining iterations do nothing; the one return at the
+        # end hands back the outputs as the early return left them.
+        return [
+            ast.Assign(
+                targets=[ast.Name(id="_ret", ctx=ast.Store())],
+                value=_jnp("bool_", [ast.Constant(False)]),
+            ),
+            *_guard_after_returns(body[:-1]),
+            final,
+        ]
 
     class Returns(ast.NodeTransformer):
         def visit_Return(self, node: ast.Return) -> Any:
diff --git a/tests/test_jax_transform.py b/tests/test_jax_transform.py
index a251fac..b461f76 100644
--- a/tests/test_jax_transform.py
+++ b/tests/test_jax_transform.py
@@ -800,3 +800,80 @@ def run(bad, worse):
         sys.path.remove(str(out))
         for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"):
             sys.modules.pop(f"early_mod{suffix}", None)
+
+
+RETURNS_IN_A_LOOP = """\
+module loopret_mod
+  implicit none
+  type coefs_type
+    real(8), allocatable :: coef(:, :)
+  end type coefs_type
+contains
+  subroutine init_coefs( nz, ngrdcol, c )
+    integer, intent(in) :: nz, ngrdcol
+    type(coefs_type), intent(inout) :: c
+    allocate( c%coef(1:ngrdcol, 1:nz) )
+    c%coef = 2.0d0
+  end subroutine init_coefs
+  subroutine apply( nzt, ngrdcol, c, x, y, err )
+    integer, intent(in) :: nzt, ngrdcol
+    type(coefs_type), intent(in) :: c
+    real(8), intent(inout), dimension(ngrdcol, nzt) :: x
+    real(8), intent(out), dimension(ngrdcol, nzt) :: y
+    integer, intent(in) :: err(ngrdcol)
+    integer :: i
+    x = x * c%coef(:, 1:nzt)
+    y = x
+    do i = 1, ngrdcol
+      if ( err(i) /= 0 ) then
+        return
+      end if
+      y(i, :) = y(i, :) + 1.0d0
+    end do
+    y = y + 10.0d0
+  end subroutine apply
+end module loopret_mod
+"""
+
+
+def test_an_early_return_inside_a_loop_becomes_a_flag(tmp_path: Path) -> None:
+    """advance_clubb_core checks the error code per column inside a loop and
+    returns. No branch structure holds that; the return sets a flag, every
+    later statement runs under it, and the remaining iterations do nothing."""
+    import importlib
+    import sys
+
+    from recast.transform.jax.tree import TreeToJax
+    from recast.transform.numpy.tree import TreeConventions
+
+    (tmp_path / "loopret_mod.f90").write_text(RETURNS_IN_A_LOOP)
+    frontend = FortranFrontend(flatten={"patch_count": "ngrdcol", "bounds_pattern": r"^ngrdcol$"})
+    unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:loopret_mod")
+    facts = frontend.analyze(unit, tmp_path)
+    candidate = TreeToJax(TreeConventions()).apply(unit, facts, {"root": str(tmp_path)})
+    assert "apply_flat" in candidate.notes["jax"]["kernels"], candidate.notes["jax"]
+    out = tmp_path / "emitted"
+    out.mkdir()
+    for path, content in candidate.files.items():
+        (out / path.name).write_bytes(content)
+    sys.path.insert(0, str(out))
+    try:
+        module = importlib.import_module("loopret_mod_jax")
+        import numpy as np
+
+        coef = np.full((2, 1), 2.0, order="F")
+
+        def run(err):
+            x = np.ones((2, 1), order="F")
+            _x, y = module.apply_flat(
+                1, 2, x, np.zeros((2, 1), order="F"), np.array(err, dtype=np.int32), 1, coef
+            )
+            return np.asarray(y).ravel().tolist()
+
+        assert run([0, 0]) == [13.0, 13.0]
+        assert run([0, 1]) == [3.0, 2.0]  # the second column returns before its +1 and the +10
+        assert run([1, 0]) == [2.0, 2.0]
+    finally:
+        sys.path.remove(str(out))
+        for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"):
+            sys.modules.pop(f"loopret_mod{suffix}", None)

From 9f55bd53676ade7cdd00923e1ceee2a01daf795c Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 17:38:40 -0600
Subject: [PATCH 51/73] The return flag's guards: decide before the rewrite
 takes the returns away

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/jax/tree.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/recast/transform/jax/tree.py b/src/recast/transform/jax/tree.py
index 1dfd92a..af28338 100644
--- a/src/recast/transform/jax/tree.py
+++ b/src/recast/transform/jax/tree.py
@@ -1636,6 +1636,7 @@ def _guard_after_returns(stmts: list[ast.stmt]) -> list[ast.stmt]:
     out: list[ast.stmt] = []
     exited = False
     for statement in stmts:
+        had_return = _has_return([statement])  # before the rewrite takes the returns away
         if isinstance(statement, ast.Return):
             rewritten: ast.stmt = ast.copy_location(
                 ast.Assign(
@@ -1657,7 +1658,7 @@ def _guard_after_returns(stmts: list[ast.stmt]) -> list[ast.stmt]:
                 orelse=[],
             )
         out.append(rewritten)
-        if _has_return([statement]):
+        if had_return:
             exited = True
     return out
 

From 7f9b30c12f8edfba88d980c14de263232d29fdd1 Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 18:09:02 -0600
Subject: [PATCH 52/73] JAX port: the anchor's unpack of an elided call buffer
 is bound, not read

advance_clubb_core's anchor unpacks what a solver hands back -- stats =
_out[0], pdf_params = _out[4] -- from the buffer the flat rewrite
elides at the call. The flat outputs bound to the actuals, so an object's
unpack is already true and goes; an array's is the actual it names.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/jax/tree.py | 27 +++++++++++++
 tests/test_jax_transform.py      | 66 ++++++++++++++++++++++++++++++++
 2 files changed, 93 insertions(+)

diff --git a/src/recast/transform/jax/tree.py b/src/recast/transform/jax/tree.py
index af28338..4e6a64c 100644
--- a/src/recast/transform/jax/tree.py
+++ b/src/recast/transform/jax/tree.py
@@ -354,6 +354,33 @@ def visit_Assign(self, node: ast.Assign) -> Any:
             return None
         if isinstance(node.value, ast.Call) and self._flat_callee(node.value) is not None:
             return self._rewrite_call(node.targets[0], node.value)
+        value = node.value
+        if (
+            len(node.targets) == 1
+            and isinstance(value, ast.Subscript)
+            and isinstance(value.value, ast.Name)
+            and value.value.id in self.buffer_outs
+            and isinstance(value.slice, ast.Constant)
+            and isinstance(value.slice.value, int)
+        ):
+            # ``stats = _out[0]`` after ``_out = callee(...)`` (the anchor
+            # unpacking an object or an array the callee handed back): the
+            # buffer was elided at the call and the flat outputs bound to
+            # the actuals directly, so an object's unpack is already true
+            # and an array's is the actual it names -- itself, usually.
+            actuals = self.buffer_outs[value.value.id]
+            at = value.slice.value
+            if at >= len(actuals):
+                raise NotFlat(f"{ast.unparse(node)}: the elided buffer has no slot {at}")
+            target = node.targets[0]
+            if ast.unparse(target) == actuals[at]:
+                return None
+            if isinstance(target, ast.Name) and self.spelling.object_of(target) is not None:
+                return None
+            bound_value = ast.parse(actuals[at], mode="eval").body
+            return self.visit(
+                ast.copy_location(ast.Assign(targets=[target], value=bound_value), node)
+            )
         pending = self._companion_writes(node.value) if isinstance(node.value, ast.Call) else []
         self.generic_visit(node)
         if pending:
diff --git a/tests/test_jax_transform.py b/tests/test_jax_transform.py
index b461f76..c292357 100644
--- a/tests/test_jax_transform.py
+++ b/tests/test_jax_transform.py
@@ -877,3 +877,69 @@ def run(err):
         sys.path.remove(str(out))
         for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"):
             sys.modules.pop(f"loopret_mod{suffix}", None)
+
+
+UNPACKS_AN_OBJECT = """\
+module unpack_mod
+  implicit none
+  type knobs_type
+    real(8) :: gain = 2.0d0
+  end type knobs_type
+contains
+  subroutine scale( nzt, ngrdcol, k, x, y )
+    integer, intent(in) :: nzt, ngrdcol
+    type(knobs_type), intent(inout) :: k
+    real(8), intent(inout), dimension(ngrdcol, nzt) :: x
+    real(8), intent(inout), dimension(ngrdcol, nzt) :: y
+    k%gain = k%gain * 2.0d0
+    x = x * 2.0d0
+    y = x + 1.0d0
+  end subroutine scale
+  subroutine step( nzt, ngrdcol, k, x, y )
+    integer, intent(in) :: nzt, ngrdcol
+    type(knobs_type), intent(inout) :: k
+    real(8), intent(inout), dimension(ngrdcol, nzt) :: x
+    real(8), intent(inout), dimension(ngrdcol, nzt) :: y
+    call scale( nzt, ngrdcol, k, x, y )
+    y = y + k%gain
+  end subroutine step
+end module unpack_mod
+"""
+
+
+def test_unpacking_an_object_from_the_elided_call_buffer_is_already_true(tmp_path: Path) -> None:
+    """advance_clubb_core's anchor unpacks what a solver hands back --
+    ``stats = _out[0]``, ``pdf_params = _out[4]`` -- from the buffer the
+    flat rewrite elides at the call: the flat outputs bound to the actuals,
+    so the object's unpack is already true and the statement goes; an
+    array's is the actual it names."""
+    import importlib
+    import sys
+
+    from recast.transform.jax.tree import TreeToJax
+    from recast.transform.numpy.tree import TreeConventions
+
+    (tmp_path / "unpack_mod.f90").write_text(UNPACKS_AN_OBJECT)
+    frontend = FortranFrontend(flatten={"patch_count": "ngrdcol", "bounds_pattern": r"^ngrdcol$"})
+    unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:unpack_mod")
+    facts = frontend.analyze(unit, tmp_path)
+    candidate = TreeToJax(TreeConventions()).apply(unit, facts, {"root": str(tmp_path)})
+    assert "step_flat" in candidate.notes["jax"]["kernels"], candidate.notes["jax"]
+    assert "_out[0]" in candidate.files[Path("unpack_mod_numpy.py")].decode()
+    out = tmp_path / "emitted"
+    out.mkdir()
+    for path, content in candidate.files.items():
+        (out / path.name).write_bytes(content)
+    sys.path.insert(0, str(out))
+    try:
+        module = importlib.import_module("unpack_mod_jax")
+        import numpy as np
+
+        x = np.ones((1, 2), order="F")
+        result = module.step_flat(2, 1, x, np.zeros((1, 2), order="F"), np.float64(2.0))
+        got = [np.asarray(v).tolist() for v in result]
+        assert got == [[[2.0, 2.0]], [[7.0, 7.0]], 4.0]
+    finally:
+        sys.path.remove(str(out))
+        for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"):
+            sys.modules.pop(f"unpack_mod{suffix}", None)

From 836b5257ba382c0bf5631f19dc4e85a021a94f7e Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 18:34:58 -0600
Subject: [PATCH 53/73] The elided call buffer keeps one slot per output the
 callee returns

The anchor unpacks _out[k] by position; an object's outputs are bound by
component and left their slot out, so advance_clubb_core's
err_info = _out[10] pointed past the list.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/jax/tree.py | 16 +++++++++++-----
 tests/test_jax_transform.py      |  8 ++++----
 2 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/src/recast/transform/jax/tree.py b/src/recast/transform/jax/tree.py
index 4e6a64c..44592ac 100644
--- a/src/recast/transform/jax/tree.py
+++ b/src/recast/transform/jax/tree.py
@@ -193,7 +193,7 @@ def __init__(
         self.inits: dict[str, str] = {}  # local -> "int32" | "float64", from its guard init
         self.loop_vars: set[str] = set()  # loop counters: int64 under x64, cast when stored
         self.state_params: list[str] = []  # a helper's module state, taken as parameters
-        self.buffer_outs: dict[str, list[str]] = {}  # anchor _out buffers, elided
+        self.buffer_outs: dict[str, list[str | None]] = {}  # anchor _out buffers, elided
         self.masked: list[str] = []  # statements whose dynamic slices became masks
         self.static_loops: list[str] = []  # loops whose trip count became static
 
@@ -373,11 +373,14 @@ def visit_Assign(self, node: ast.Assign) -> Any:
             if at >= len(actuals):
                 raise NotFlat(f"{ast.unparse(node)}: the elided buffer has no slot {at}")
             target = node.targets[0]
-            if ast.unparse(target) == actuals[at]:
-                return None
             if isinstance(target, ast.Name) and self.spelling.object_of(target) is not None:
+                return None  # an object: its components bound at the call
+            actual = actuals[at]
+            if actual is None:
+                raise NotFlat(f"{ast.unparse(node)}: the elided buffer's slot {at} has no actual")
+            if ast.unparse(target) == actual:
                 return None
-            bound_value = ast.parse(actuals[at], mode="eval").body
+            bound_value = ast.parse(actual, mode="eval").body
             return self.visit(
                 ast.copy_location(ast.Assign(targets=[target], value=bound_value), node)
             )
@@ -1294,8 +1297,11 @@ def _rewrite_call(self, target: ast.expr | None, call: ast.Call) -> Any:
                 for name in original_outs
                 if name.lower() in actual_by_dummy
             }
+            # Positional, one entry per output the callee returns: the
+            # anchor unpacks ``_out[k]`` by that position, and an object's
+            # outputs (bound by component) leave their slot empty.
             self.buffer_outs[buffered.id] = [
-                ast.unparse(slot[name]) for name in original_outs if name in slot
+                ast.unparse(slot[name]) if name in slot else None for name in original_outs
             ]
         targets: list[ast.expr] = []
         follow: list[ast.stmt] = []
diff --git a/tests/test_jax_transform.py b/tests/test_jax_transform.py
index c292357..7522796 100644
--- a/tests/test_jax_transform.py
+++ b/tests/test_jax_transform.py
@@ -886,11 +886,11 @@ def run(err):
     real(8) :: gain = 2.0d0
   end type knobs_type
 contains
-  subroutine scale( nzt, ngrdcol, k, x, y )
+  subroutine scale( nzt, ngrdcol, x, y, k )
     integer, intent(in) :: nzt, ngrdcol
-    type(knobs_type), intent(inout) :: k
     real(8), intent(inout), dimension(ngrdcol, nzt) :: x
     real(8), intent(inout), dimension(ngrdcol, nzt) :: y
+    type(knobs_type), intent(inout) :: k
     k%gain = k%gain * 2.0d0
     x = x * 2.0d0
     y = x + 1.0d0
@@ -900,7 +900,7 @@ def run(err):
     type(knobs_type), intent(inout) :: k
     real(8), intent(inout), dimension(ngrdcol, nzt) :: x
     real(8), intent(inout), dimension(ngrdcol, nzt) :: y
-    call scale( nzt, ngrdcol, k, x, y )
+    call scale( nzt, ngrdcol, x, y, k )
     y = y + k%gain
   end subroutine step
 end module unpack_mod
@@ -925,7 +925,7 @@ def test_unpacking_an_object_from_the_elided_call_buffer_is_already_true(tmp_pat
     facts = frontend.analyze(unit, tmp_path)
     candidate = TreeToJax(TreeConventions()).apply(unit, facts, {"root": str(tmp_path)})
     assert "step_flat" in candidate.notes["jax"]["kernels"], candidate.notes["jax"]
-    assert "_out[0]" in candidate.files[Path("unpack_mod_numpy.py")].decode()
+    assert "k = _out[2]" in candidate.files[Path("unpack_mod_numpy.py")].decode()
     out = tmp_path / "emitted"
     out.mkdir()
     for path, content in candidate.files.items():

From 98e6392e6346abef3a9a14c162490104accc37cc Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 19:30:42 -0600
Subject: [PATCH 54/73] JAX port: a companion procedure the port left on the
 host stays under its guard

advance_clubb_core calls numerical_check's parameterization check under
clubb_at_least_debug_level_api(2), false with statistics off; the check
takes character arguments and its port leaves it on the host. The call
stays the host's under its trace-time guard (a stand-in's function of
constants is a static test), is never traced while the guard holds, fails
by name if it ever does, and the note names it. Refusing the whole step
for a check that never runs was the alternative.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/jax/backend.py | 15 +++++-
 src/recast/transform/jax/tree.py    | 20 +++++++-
 tests/test_jax_transform.py         | 77 +++++++++++++++++++++++++++++
 3 files changed, 110 insertions(+), 2 deletions(-)

diff --git a/src/recast/transform/jax/backend.py b/src/recast/transform/jax/backend.py
index f598065..e3e7b6a 100644
--- a/src/recast/transform/jax/backend.py
+++ b/src/recast/transform/jax/backend.py
@@ -403,6 +403,17 @@ def constant(node):
             return constant(node.left) and constant(node.right)
         if isinstance(node, ast.UnaryOp):
             return constant(node.operand)
+        if (
+            isinstance(node, ast.Call)
+            and isinstance(node.func, ast.Attribute)
+            and isinstance(node.func.value, ast.Name)
+            and all(constant(a) for a in node.args)
+            and not node.keywords
+        ):
+            # ``_error_code.clubb_at_least_debug_level_api(2)``: a stand-in's
+            # function of constants, a Python value at trace time. (A kernel
+            # of constants is a concrete array; ``_f_concrete`` tells.)
+            return True
         return False
 
     # ``RUNGE_KUTTA_TYPE == I_10``: a comparison over module constants
@@ -412,7 +423,9 @@ def constant(node):
         return constant(test.left) and all(constant(c) for c in test.comparators)
     if isinstance(test, ast.BoolOp):
         return all(_static_test(v, statics) for v in test.values)
-    return False
+    if isinstance(test, ast.UnaryOp) and isinstance(test.op, ast.Not):
+        return _static_test(test.operand, statics)
+    return constant(test)
 
 
 def _names(ids, ctx):
diff --git a/src/recast/transform/jax/tree.py b/src/recast/transform/jax/tree.py
index 44592ac..3a7d3b1 100644
--- a/src/recast/transform/jax/tree.py
+++ b/src/recast/transform/jax/tree.py
@@ -194,6 +194,7 @@ def __init__(
         self.loop_vars: set[str] = set()  # loop counters: int64 under x64, cast when stored
         self.state_params: list[str] = []  # a helper's module state, taken as parameters
         self.buffer_outs: dict[str, list[str | None]] = {}  # anchor _out buffers, elided
+        self.host_calls: list[str] = []  # companion procedures left on the host, by name
         self.masked: list[str] = []  # statements whose dynamic slices became masks
         self.static_loops: list[str] = []  # loops whose trip count became static
 
@@ -901,7 +902,17 @@ def _companion_call(self, node: ast.Call) -> ast.AST:
             # answers, not physics, and are left as they are.
             return node
         if func.attr not in port["kernels"]:
-            raise NotFlat(f"calls {module}.{func.attr}, which its port did not lower")
+            # The companion's port left this one on the host (CLUBB's
+            # numerical_check.parameterization_check: character arguments,
+            # writes). The call stays the host's, under whatever guard the
+            # anchor put it -- ``if clubb_at_least_debug_level_api(2)``, a
+            # stand-in's answer at trace time, false with statistics off --
+            # and is never traced while the guard holds. If the guard ever
+            # lets it through, the trace fails on it by name, which the
+            # gate reports; refusing the whole kernel for a check that
+            # never runs was the alternative.
+            self.host_calls.append(f"{module}.{func.attr}")
+            return node
         closure: list[ast.expr] = []
         for state in port["closures"].get(func.attr, []):
             flat = f"{module}__{state}"
@@ -2683,6 +2694,7 @@ def flattened_module(
     specialized: dict[str, tuple[ast.FunctionDef, FlatPlan, list[str]]] = {}
     aborts: dict[str, list[str]] = {}
     masked: dict[str, list[str]] = {}
+    host_calls: dict[str, list[str]] = {}
     static_loops: dict[str, list[str]] = {}
     companions: set[str] = set()
     planned = {p.subprogram["name"] for p in plans}
@@ -2704,6 +2716,8 @@ def flattened_module(
             aborts[plan.name] = rewrite.aborts
         if rewrite.masked:
             masked[plan.name] = rewrite.masked
+        if rewrite.host_calls:
+            host_calls[plan.name] = sorted(set(rewrite.host_calls))
         if rewrite.static_loops:
             static_loops[plan.name] = rewrite.static_loops
         companions |= rewrite.companions
@@ -2726,6 +2740,8 @@ def flattened_module(
             aborts[name] = rewrite.aborts
         if rewrite.masked:
             masked[name] = rewrite.masked
+        if rewrite.host_calls:
+            host_calls[name] = sorted(set(rewrite.host_calls))
         if rewrite.static_loops:
             static_loops[name] = rewrite.static_loops
         if rewrite.state_params:
@@ -2751,6 +2767,7 @@ def flattened_module(
         "aborts_dropped": aborts,
         "masked": masked,
         "static_loops": static_loops,
+        "host_calls": host_calls,
         "companions": sorted(companions),
     }
     return module, {**interface, "subprograms": [*interface["subprograms"], *entries]}, notes
@@ -2851,6 +2868,7 @@ def apply(self, unit: Unit, facts: Facts, config: dict[str, Any]) -> Candidate:
                     "aborts_dropped": dict(sorted(flat_notes["aborts_dropped"].items())),
                     "masked": dict(sorted(flat_notes["masked"].items())),
                     "static_loops": dict(sorted(flat_notes["static_loops"].items())),
+                    "host_calls": dict(sorted(flat_notes.get("host_calls", {}).items())),
                     "companions": sorted(ported),
                     "runtime": f"{runtime_stem}.py",
                     "_ports": ported,
diff --git a/tests/test_jax_transform.py b/tests/test_jax_transform.py
index 7522796..d1f3b53 100644
--- a/tests/test_jax_transform.py
+++ b/tests/test_jax_transform.py
@@ -943,3 +943,80 @@ def test_unpacking_an_object_from_the_elided_call_buffer_is_already_true(tmp_pat
         sys.path.remove(str(out))
         for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"):
             sys.modules.pop(f"unpack_mod{suffix}", None)
+
+
+CHECKS = """\
+module checks_mod
+  implicit none
+contains
+  subroutine complain( n, x, what )
+    integer, intent(in) :: n
+    real(8), intent(in) :: x(n)
+    character(len=*), intent(in) :: what
+    integer :: i
+    i = 1
+    do while ( i <= n )
+      if ( x(i) < 0.0d0 ) print *, what, i
+      i = i + 1
+    end do
+  end subroutine complain
+end module checks_mod
+"""
+
+GUARDED_CHECK = """\
+module guarded_mod
+  use checks_mod, only: complain
+  implicit none
+  integer, parameter :: debug_level = 0
+contains
+  subroutine step( n, x, y )
+    integer, intent(in) :: n
+    real(8), intent(in) :: x(n)
+    real(8), intent(out) :: y(n)
+    y = 2.0d0 * x
+    if ( debug_level >= 2 ) then
+      call complain( n, y, "negative y" )
+    end if
+  end subroutine step
+end module guarded_mod
+"""
+
+
+def test_a_companion_procedure_the_port_left_on_the_host_stays_under_its_guard(
+    tmp_path: Path,
+) -> None:
+    """CLUBB's advance_clubb_core calls numerical_check's parameterization
+    check under ``clubb_at_least_debug_level_api(2)``, false with statistics
+    off; the check takes character arguments and its port leaves it on the
+    host. Refusing the whole step for a call that never runs was the
+    alternative: the call stays the host's under its trace-time guard, and
+    the note names it."""
+    import importlib
+    import sys
+
+    from recast.transform.jax.tree import TreeToJax
+    from recast.transform.numpy.tree import TreeConventions
+
+    (tmp_path / "checks_mod.f90").write_text(CHECKS)
+    (tmp_path / "guarded_mod.f90").write_text(GUARDED_CHECK)
+    frontend = FortranFrontend(flatten=True)
+    unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:guarded_mod")
+    facts = frontend.analyze(unit, tmp_path)
+    candidate = TreeToJax(TreeConventions()).apply(unit, facts, {"root": str(tmp_path)})
+    assert "step" in candidate.notes["jax"]["kernels"], candidate.notes["jax"]
+    assert candidate.notes["jax"]["host_calls"] == {"step": ["checks_mod.complain"]}
+    out = tmp_path / "emitted"
+    out.mkdir()
+    for path, content in candidate.files.items():
+        (out / path.name).write_bytes(content)
+    sys.path.insert(0, str(out))
+    try:
+        module = importlib.import_module("guarded_mod_jax")
+        import numpy as np
+
+        assert np.asarray(module.step(2, np.array([1.0, -1.0]))).tolist() == [2.0, -2.0]
+    finally:
+        sys.path.remove(str(out))
+        for name in list(sys.modules):
+            if name.startswith(("guarded_mod", "checks_mod")):
+                sys.modules.pop(name, None)

From 9f0f019bf148757d2c62ee09f9f179339d74e42e Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 19:32:05 -0600
Subject: [PATCH 55/73] The guarded host-call test guards on a static argument,
 as CLUBB's stand-in does

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 tests/test_jax_transform.py | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/tests/test_jax_transform.py b/tests/test_jax_transform.py
index d1f3b53..e4b718d 100644
--- a/tests/test_jax_transform.py
+++ b/tests/test_jax_transform.py
@@ -967,10 +967,9 @@ def test_unpacking_an_object_from_the_elided_call_buffer_is_already_true(tmp_pat
 module guarded_mod
   use checks_mod, only: complain
   implicit none
-  integer, parameter :: debug_level = 0
 contains
-  subroutine step( n, x, y )
-    integer, intent(in) :: n
+  subroutine step( n, debug_level, x, y )
+    integer, intent(in) :: n, debug_level
     real(8), intent(in) :: x(n)
     real(8), intent(out) :: y(n)
     y = 2.0d0 * x
@@ -1014,7 +1013,7 @@ def test_a_companion_procedure_the_port_left_on_the_host_stays_under_its_guard(
         module = importlib.import_module("guarded_mod_jax")
         import numpy as np
 
-        assert np.asarray(module.step(2, np.array([1.0, -1.0]))).tolist() == [2.0, -2.0]
+        assert np.asarray(module.step(2, 0, np.array([1.0, -1.0]))).tolist() == [2.0, -2.0]
     finally:
         sys.path.remove(str(out))
         for name in list(sys.modules):

From ea1a842efa43853efdf136c97499f31da71ff1f0 Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 19:34:04 -0600
Subject: [PATCH 56/73] JAX port: a call statement is not an inert log line

The branch pruning took any expression statement for a log line the
anchor had left as a string and dropped the branch; a bare call -- a
procedure the port left on the host -- went with it, silently. Only
pass and a bare constant are inert.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/jax/tree.py | 16 +++++++++++++---
 tests/test_jax_transform.py      | 15 ++++++---------
 2 files changed, 19 insertions(+), 12 deletions(-)

diff --git a/src/recast/transform/jax/tree.py b/src/recast/transform/jax/tree.py
index 3a7d3b1..8faedc9 100644
--- a/src/recast/transform/jax/tree.py
+++ b/src/recast/transform/jax/tree.py
@@ -660,7 +660,7 @@ def visit_If(self, node: ast.If) -> Any:
         # under tracing, and the recorded run it is gated on never did; the
         # check is dropped and named on the candidate, so the evidence says
         # the kernel no longer aborts where the model would.
-        if node.body and all(isinstance(s, (ast.Raise, ast.Expr, ast.Pass)) for s in node.body):
+        if node.body and all(isinstance(s, ast.Raise) or _inert(s) for s in node.body):
             if any(isinstance(s, ast.Raise) for s in node.body):
                 self.aborts.append(ast.unparse(node.test))
                 if not node.orelse:
@@ -673,9 +673,11 @@ def visit_If(self, node: ast.If) -> Any:
                     replaced.extend(seen if isinstance(seen, list) else [seen])
                 return replaced or None
         self.generic_visit(node)
-        if all(isinstance(s, (ast.Pass, ast.Expr)) for s in node.body):
+        if all(_inert(s) for s in node.body):
             # ``if cond: write(iulog, ...)`` -- a log line the anchor already
-            # left as ``pass``; nothing to carry. An else branch keeps its
+            # left as ``pass``; nothing to carry. A call statement is not
+            # inert: a procedure the port left on the host, under its guard,
+            # must stay where it is. An else branch keeps its
             # guard, inverted: the fold of an early return leaves exactly
             # ``if returned: pass else: <the rest>``, and running the rest
             # unconditionally was the returned path's outputs overwritten.
@@ -1636,6 +1638,14 @@ def _dim_sources(args: list[dict[str, Any]]) -> dict[str, tuple[str, int]]:
     return sources
 
 
+def _inert(statement: ast.stmt) -> bool:
+    """A statement with nothing to carry: ``pass``, or a bare constant (a
+    docstring, a log line the anchor left as a string)."""
+    return isinstance(statement, ast.Pass) or (
+        isinstance(statement, ast.Expr) and isinstance(statement.value, ast.Constant)
+    )
+
+
 def _has_return(stmts: list[ast.stmt]) -> bool:
     return any(isinstance(n, ast.Return) for s in stmts for n in ast.walk(s))
 
diff --git a/tests/test_jax_transform.py b/tests/test_jax_transform.py
index e4b718d..4a59928 100644
--- a/tests/test_jax_transform.py
+++ b/tests/test_jax_transform.py
@@ -949,16 +949,12 @@ def test_unpacking_an_object_from_the_elided_call_buffer_is_already_true(tmp_pat
 module checks_mod
   implicit none
 contains
-  subroutine complain( n, x, what )
+  subroutine complain( n, x, msg )
     integer, intent(in) :: n
     real(8), intent(in) :: x(n)
-    character(len=*), intent(in) :: what
-    integer :: i
-    i = 1
-    do while ( i <= n )
-      if ( x(i) < 0.0d0 ) print *, what, i
-      i = i + 1
-    end do
+    character(len=*), intent(out) :: msg
+    msg = "fine"
+    if ( any( x < 0.0d0 ) ) msg = "negative"
   end subroutine complain
 end module checks_mod
 """
@@ -972,9 +968,10 @@ def test_unpacking_an_object_from_the_elided_call_buffer_is_already_true(tmp_pat
     integer, intent(in) :: n, debug_level
     real(8), intent(in) :: x(n)
     real(8), intent(out) :: y(n)
+    character(len=16) :: msg
     y = 2.0d0 * x
     if ( debug_level >= 2 ) then
-      call complain( n, y, "negative y" )
+      call complain( n, y, msg )
     end if
   end subroutine step
 end module guarded_mod

From 078256cfe79c77854ff4d5565eaff210ece322c5 Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 20:22:30 -0600
Subject: [PATCH 57/73] JAX port: a flat companion the port left on the host
 stays under its guard too

advance_clubb_core's parameterization check takes gr and err_info, so the
call reaches the flat-callee rewrite, which refused when the companion's
port had no kernel for it; the same host-call policy applies there.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/jax/tree.py | 20 ++++++++-
 tests/test_jax_transform.py      | 75 ++++++++++++++++++++++++++++++++
 2 files changed, 93 insertions(+), 2 deletions(-)

diff --git a/src/recast/transform/jax/tree.py b/src/recast/transform/jax/tree.py
index 8faedc9..e29a187 100644
--- a/src/recast/transform/jax/tree.py
+++ b/src/recast/transform/jax/tree.py
@@ -1194,6 +1194,24 @@ def visit_Name(self, node: ast.Name) -> ast.AST:
     def _rewrite_call(self, target: ast.expr | None, call: ast.Call) -> Any:
         callee = self._flat_callee(call)
         assert callee is not None and self.plan is not None
+        if isinstance(call.func, ast.Attribute) and isinstance(call.func.value, ast.Name):
+            module = self.spelling.modules.get(call.func.value.id)
+            port = self.ports.get(module) if module is not None else None
+            if port is not None and callee.name not in port["kernels"]:
+                # The companion's port left this flat function on the host
+                # (CLUBB's numerical_check.parameterization_check: character
+                # arguments). The anchor's statement stays as it is, under
+                # whatever guard the anchor put it -- never traced while
+                # the guard holds, failing by name if it ever does -- and
+                # the note names it; refusing the whole kernel for a check
+                # that never runs was the alternative.
+                self.host_calls.append(f"{module}.{callee.name}")
+                kept: ast.stmt = (
+                    ast.Expr(value=call)
+                    if target is None
+                    else ast.Assign(targets=[target], value=call)
+                )
+                return ast.copy_location(kept, call)
         dummies = [a["name"] for a in callee.subprogram["args"]]
         actual_by_dummy: dict[str, ast.expr] = {}
         for at, given in enumerate(call.args):
@@ -1256,8 +1274,6 @@ def _rewrite_call(self, target: ast.expr | None, call: ast.Call) -> Any:
         elif isinstance(call.func, ast.Attribute) and isinstance(call.func.value, ast.Name):
             # A companion's flat function: its port's kernel, or nothing.
             module = self.spelling.modules[call.func.value.id]
-            if callee.name not in self.ports[module]["kernels"]:
-                raise NotFlat(f"calls {module}.{callee.name}, which its port did not lower")
             self.companions.add(call.func.value.id)
             func = ast.Attribute(
                 value=ast.Name(id=f"{call.func.value.id}_jax", ctx=ast.Load()),
diff --git a/tests/test_jax_transform.py b/tests/test_jax_transform.py
index 4a59928..2582554 100644
--- a/tests/test_jax_transform.py
+++ b/tests/test_jax_transform.py
@@ -1016,3 +1016,78 @@ def test_a_companion_procedure_the_port_left_on_the_host_stays_under_its_guard(
         for name in list(sys.modules):
             if name.startswith(("guarded_mod", "checks_mod")):
                 sys.modules.pop(name, None)
+
+
+OBJECT_CHECK = """\
+module ocheck_mod
+  implicit none
+  type knobs_type
+    real(8) :: gain = 2.0d0
+  end type knobs_type
+contains
+  subroutine inspect( n, k, x, msg )
+    integer, intent(in) :: n
+    type(knobs_type), intent(in) :: k
+    real(8), intent(in) :: x(n)
+    character(len=*), intent(out) :: msg
+    msg = "fine"
+    if ( any( x < k%gain ) ) msg = "small"
+  end subroutine inspect
+end module ocheck_mod
+"""
+
+GUARDED_OBJECT_CHECK = """\
+module oguarded_mod
+  use ocheck_mod, only: knobs_type, inspect
+  implicit none
+contains
+  subroutine step( n, debug_level, k, x, y )
+    integer, intent(in) :: n, debug_level
+    type(knobs_type), intent(in) :: k
+    real(8), intent(in) :: x(n)
+    real(8), intent(out) :: y(n)
+    character(len=16) :: msg
+    y = k%gain * x
+    if ( debug_level >= 2 ) then
+      call inspect( n, k, y, msg )
+    end if
+  end subroutine step
+end module oguarded_mod
+"""
+
+
+def test_a_flat_companion_the_port_left_on_the_host_stays_under_its_guard(tmp_path: Path) -> None:
+    """The same, through the flat-callee path: the check takes the object
+    (advance_clubb_core's parameterization check takes gr and err_info), so
+    it has a plan, and the rewrite of the call into the port's flat kernel
+    found none."""
+    import importlib
+    import sys
+
+    from recast.transform.jax.tree import TreeToJax
+    from recast.transform.numpy.tree import TreeConventions
+
+    (tmp_path / "ocheck_mod.f90").write_text(OBJECT_CHECK)
+    (tmp_path / "oguarded_mod.f90").write_text(GUARDED_OBJECT_CHECK)
+    frontend = FortranFrontend(flatten=True)
+    unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:oguarded_mod")
+    facts = frontend.analyze(unit, tmp_path)
+    candidate = TreeToJax(TreeConventions()).apply(unit, facts, {"root": str(tmp_path)})
+    assert "step_flat" in candidate.notes["jax"]["kernels"], candidate.notes["jax"]
+    assert candidate.notes["jax"]["host_calls"] == {"step_flat": ["ocheck_mod.inspect"]}
+    out = tmp_path / "emitted"
+    out.mkdir()
+    for path, content in candidate.files.items():
+        (out / path.name).write_bytes(content)
+    sys.path.insert(0, str(out))
+    try:
+        module = importlib.import_module("oguarded_mod_jax")
+        import numpy as np
+
+        got = module.step_flat(2, 0, np.array([1.0, -1.0]), np.float64(2.0))
+        assert np.asarray(got).tolist() == [2.0, -2.0]
+    finally:
+        sys.path.remove(str(out))
+        for name in list(sys.modules):
+            if name.startswith(("oguarded_mod", "ocheck_mod")):
+                sys.modules.pop(name, None)

From 2618467abfe89bf95ec088c3f4744db0c81e31fa Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 20:23:49 -0600
Subject: [PATCH 58/73] Call the object-check fixture's flat kernel by its
 signature

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 tests/test_jax_transform.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_jax_transform.py b/tests/test_jax_transform.py
index 2582554..4760704 100644
--- a/tests/test_jax_transform.py
+++ b/tests/test_jax_transform.py
@@ -1084,7 +1084,7 @@ def test_a_flat_companion_the_port_left_on_the_host_stays_under_its_guard(tmp_pa
         module = importlib.import_module("oguarded_mod_jax")
         import numpy as np
 
-        got = module.step_flat(2, 0, np.array([1.0, -1.0]), np.float64(2.0))
+        got = module.step_flat(2, 0, np.array([1.0, -1.0]), np.zeros(2), 1, np.float64(2.0))
         assert np.asarray(got).tolist() == [2.0, -2.0]
     finally:
         sys.path.remove(str(out))

From 675ece968225e42829f872f8183419e1c246b01d Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 21:09:48 -0600
Subject: [PATCH 59/73] A host-kept call releases the anchor's result buffer to
 the host

The unpacks that follow a call the port left on the host read the
anchor's _out under the same guard; the elided-buffer entry of an earlier
call must not answer for them.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/jax/tree.py | 4 ++++
 tests/test_jax_transform.py      | 9 ++++++---
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/src/recast/transform/jax/tree.py b/src/recast/transform/jax/tree.py
index e29a187..8850d2b 100644
--- a/src/recast/transform/jax/tree.py
+++ b/src/recast/transform/jax/tree.py
@@ -1206,6 +1206,10 @@ def _rewrite_call(self, target: ast.expr | None, call: ast.Call) -> Any:
                 # the note names it; refusing the whole kernel for a check
                 # that never runs was the alternative.
                 self.host_calls.append(f"{module}.{callee.name}")
+                if isinstance(target, ast.Name):
+                    # The anchor's result buffer is the host's now; the
+                    # unpacks that follow read it, under the same guard.
+                    self.buffer_outs.pop(target.id, None)
                 kept: ast.stmt = (
                     ast.Expr(value=call)
                     if target is None
diff --git a/tests/test_jax_transform.py b/tests/test_jax_transform.py
index 4760704..5f3e3da 100644
--- a/tests/test_jax_transform.py
+++ b/tests/test_jax_transform.py
@@ -1027,11 +1027,12 @@ def test_a_companion_procedure_the_port_left_on_the_host_stays_under_its_guard(
 contains
   subroutine inspect( n, k, x, msg )
     integer, intent(in) :: n
-    type(knobs_type), intent(in) :: k
+    type(knobs_type), intent(inout) :: k
     real(8), intent(in) :: x(n)
     character(len=*), intent(out) :: msg
     msg = "fine"
     if ( any( x < k%gain ) ) msg = "small"
+    k%gain = k%gain + 1.0d0
   end subroutine inspect
 end module ocheck_mod
 """
@@ -1043,7 +1044,7 @@ def test_a_companion_procedure_the_port_left_on_the_host_stays_under_its_guard(
 contains
   subroutine step( n, debug_level, k, x, y )
     integer, intent(in) :: n, debug_level
-    type(knobs_type), intent(in) :: k
+    type(knobs_type), intent(inout) :: k
     real(8), intent(in) :: x(n)
     real(8), intent(out) :: y(n)
     character(len=16) :: msg
@@ -1051,6 +1052,7 @@ def test_a_companion_procedure_the_port_left_on_the_host_stays_under_its_guard(
     if ( debug_level >= 2 ) then
       call inspect( n, k, y, msg )
     end if
+    y = y + k%gain
   end subroutine step
 end module oguarded_mod
 """
@@ -1085,7 +1087,8 @@ def test_a_flat_companion_the_port_left_on_the_host_stays_under_its_guard(tmp_pa
         import numpy as np
 
         got = module.step_flat(2, 0, np.array([1.0, -1.0]), np.zeros(2), 1, np.float64(2.0))
-        assert np.asarray(got).tolist() == [2.0, -2.0]
+        y, gain = (np.asarray(v).tolist() for v in got)
+        assert y == [4.0, 0.0] and gain == 2.0
     finally:
         sys.path.remove(str(out))
         for name in list(sys.modules):

From ebaecc2dcf5976adfc038e3eaa8d3a5452043c08 Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 21:58:29 -0600
Subject: [PATCH 60/73] The elided call buffer has a slot for every OUT/INOUT
 dummy, optional ones included

The anchor's return tuple has one per dummy (an absent optional's is _)
and the unpacks index it by position; CLUBB's stats = _out[3] after
calc_brunt_vaisala_freq_sqd names an optional INOUT object the list had
left out.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/jax/tree.py | 9 ++++++---
 tests/test_jax_transform.py      | 8 +++++---
 2 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/src/recast/transform/jax/tree.py b/src/recast/transform/jax/tree.py
index 8850d2b..5c48199 100644
--- a/src/recast/transform/jax/tree.py
+++ b/src/recast/transform/jax/tree.py
@@ -1299,10 +1299,13 @@ def _rewrite_call(self, target: ast.expr | None, call: ast.Call) -> Any:
             anchor_targets = list(target.elts)
         elif target is not None:
             anchor_targets = [target]
+        # Every OUT/INOUT dummy, the optional ones included: the anchor's
+        # return tuple has a slot for each (an absent optional's is ``_``),
+        # and the unpacks that follow index it by position -- CLUBB's
+        # ``stats = _out[3]`` after calc_brunt_vaisala_freq_sqd, whose
+        # ``stats`` is an optional INOUT object.
         original_outs = [
-            _py(a["name"])
-            for a in callee.subprogram["args"]
-            if a["intent"] in ("OUT", "INOUT") and not a.get("optional")
+            _py(a["name"]) for a in callee.subprogram["args"] if a["intent"] in ("OUT", "INOUT")
         ]
         if callee.subprogram["kind"] == "function":
             original_outs = ["_result", *original_outs]
diff --git a/tests/test_jax_transform.py b/tests/test_jax_transform.py
index 5f3e3da..9b541dc 100644
--- a/tests/test_jax_transform.py
+++ b/tests/test_jax_transform.py
@@ -890,8 +890,8 @@ def run(err):
     integer, intent(in) :: nzt, ngrdcol
     real(8), intent(inout), dimension(ngrdcol, nzt) :: x
     real(8), intent(inout), dimension(ngrdcol, nzt) :: y
-    type(knobs_type), intent(inout) :: k
-    k%gain = k%gain * 2.0d0
+    type(knobs_type), intent(inout), optional :: k
+    if ( present( k ) ) k%gain = k%gain * 2.0d0
     x = x * 2.0d0
     y = x + 1.0d0
   end subroutine scale
@@ -912,7 +912,9 @@ def test_unpacking_an_object_from_the_elided_call_buffer_is_already_true(tmp_pat
     ``stats = _out[0]``, ``pdf_params = _out[4]`` -- from the buffer the
     flat rewrite elides at the call: the flat outputs bound to the actuals,
     so the object's unpack is already true and the statement goes; an
-    array's is the actual it names."""
+    array's is the actual it names. The object here is an *optional* INOUT
+    (calc_brunt_vaisala_freq_sqd's ``stats``): its slot is in the anchor's
+    tuple all the same, and was not in the rewrite's list."""
     import importlib
     import sys
 

From 0be7d7199f5d6ac7c95e511221f2e2b88892b583 Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 21:59:40 -0600
Subject: [PATCH 61/73] The optional-object unpack test expects the flat rule:
 an optional dummy is absent

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 tests/test_jax_transform.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/tests/test_jax_transform.py b/tests/test_jax_transform.py
index 9b541dc..175d9d8 100644
--- a/tests/test_jax_transform.py
+++ b/tests/test_jax_transform.py
@@ -940,7 +940,11 @@ def test_unpacking_an_object_from_the_elided_call_buffer_is_already_true(tmp_pat
         x = np.ones((1, 2), order="F")
         result = module.step_flat(2, 1, x, np.zeros((1, 2), order="F"), np.float64(2.0))
         got = [np.asarray(v).tolist() for v in result]
-        assert got == [[[2.0, 2.0]], [[7.0, 7.0]], 4.0]
+        # The optional object is absent in the flat world (the plan leaves
+        # optional dummies out, as the adapter calls without them), so the
+        # callee's ``present(k)`` branch does not run: gain stays 2, y is
+        # x * 2 + 1 + 2.
+        assert got == [[[2.0, 2.0]], [[5.0, 5.0]], 2.0]
     finally:
         sys.path.remove(str(out))
         for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"):

From 4f3afbf2eb052323bc53170746bebbe98867d68c Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Fri, 4 Sep 2026 22:49:24 -0600
Subject: [PATCH 62/73] JAX lowering: a branch of bare call statements carries
 nothing

A check the port left on the host, or a kernel call whose result nothing
binds, under a guard on the data: the branch binds nothing and nothing in
it could run under a tracer, so the lowering carries nothing rather than
refusing the kernel (advance_clubb_core's parameterization check under
any(err_code == fatal)).

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/jax/backend.py |  9 +++-
 tests/test_jax_transform.py         | 64 +++++++++++++++++++++++++++++
 2 files changed, 71 insertions(+), 2 deletions(-)

diff --git a/src/recast/transform/jax/backend.py b/src/recast/transform/jax/backend.py
index e3e7b6a..8c8049f 100644
--- a/src/recast/transform/jax/backend.py
+++ b/src/recast/transform/jax/backend.py
@@ -782,13 +782,18 @@ def _cond_form(self, s, body, orelse):
             if n not in carried:
                 carried.append(n)
         if not carried:
+            # A guard around dropped logs/aborts, or around a call statement
+            # whose result nothing binds (a check the port left on the host,
+            # under a traced guard -- CLUBB's parameterization check under
+            # ``any(err_code == fatal)``): nothing to carry, and nothing a
+            # tracer could run. The tree's notes name the host calls.
             trivial = all(
                 isinstance(st, ast.Pass)
-                or (isinstance(st, ast.Expr) and isinstance(st.value, ast.Constant))
+                or (isinstance(st, ast.Expr) and isinstance(st.value, ast.Constant | ast.Call))
                 for st in [*body, *orelse]
             )
             if trivial:
-                return []  # a guard around dropped logs/aborts: nothing to carry
+                return []
             raise JaxQueue("IF with no carried effects")
         self.n += 1
         t_name, f_name = f"_true_{self.n}", f"_false_{self.n}"
diff --git a/tests/test_jax_transform.py b/tests/test_jax_transform.py
index 175d9d8..42bb095 100644
--- a/tests/test_jax_transform.py
+++ b/tests/test_jax_transform.py
@@ -1100,3 +1100,67 @@ def test_a_flat_companion_the_port_left_on_the_host_stays_under_its_guard(tmp_pa
         for name in list(sys.modules):
             if name.startswith(("oguarded_mod", "ocheck_mod")):
                 sys.modules.pop(name, None)
+
+
+SILENT_CHECK = """\
+module silent_mod
+  implicit none
+contains
+  subroutine complain( n, x )
+    integer, intent(in) :: n
+    real(8), intent(in) :: x(n)
+    if ( any( x < 0.0d0 ) ) print *, "negative", n
+  end subroutine complain
+end module silent_mod
+"""
+
+TRACED_GUARD_CHECK = """\
+module tguarded_mod
+  use silent_mod, only: complain
+  implicit none
+contains
+  subroutine step( n, x, y )
+    integer, intent(in) :: n
+    real(8), intent(in) :: x(n)
+    real(8), intent(out) :: y(n)
+    y = 2.0d0 * x
+    if ( y(1) < 0.0d0 ) then
+      call complain( n, y )
+    end if
+  end subroutine step
+end module tguarded_mod
+"""
+
+
+def test_a_host_only_call_under_a_traced_guard_carries_nothing(tmp_path: Path) -> None:
+    """The same check under a guard on the data (``any(err_code == fatal)``):
+    the branch binds nothing and nothing in it could run under a tracer, so
+    the lowering carries nothing rather than refusing the kernel."""
+    import importlib
+    import sys
+
+    from recast.transform.jax.tree import TreeToJax
+    from recast.transform.numpy.tree import TreeConventions
+
+    (tmp_path / "silent_mod.f90").write_text(SILENT_CHECK)
+    (tmp_path / "tguarded_mod.f90").write_text(TRACED_GUARD_CHECK)
+    frontend = FortranFrontend(flatten=True)
+    unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:tguarded_mod")
+    facts = frontend.analyze(unit, tmp_path)
+    candidate = TreeToJax(TreeConventions()).apply(unit, facts, {"root": str(tmp_path)})
+    assert "step" in candidate.notes["jax"]["kernels"], candidate.notes["jax"]
+    out = tmp_path / "emitted"
+    out.mkdir()
+    for path, content in candidate.files.items():
+        (out / path.name).write_bytes(content)
+    sys.path.insert(0, str(out))
+    try:
+        module = importlib.import_module("tguarded_mod_jax")
+        import numpy as np
+
+        assert np.asarray(module.step(2, np.array([-1.0, 1.0]))).tolist() == [-2.0, 2.0]
+    finally:
+        sys.path.remove(str(out))
+        for name in list(sys.modules):
+            if name.startswith(("tguarded_mod", "silent_mod")):
+                sys.modules.pop(name, None)

From b9dfc3cb451c2e2fa76bf64ba42fc19d8abde563 Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Sat, 5 Sep 2026 00:06:59 -0600
Subject: [PATCH 63/73] jax: guard the whole remainder after a return as one
 block

After a statement that may have returned, _guard_after_returns wrapped
each following statement in its own `if not _ret`. A call whose result
lands in the _out buffer (never a carry) was then separated from its
unpacks by a lax.cond with no carried effects, and the kernel fell back
to the host. Wrap the rest of the block once instead, recursively, so
the buffer and its unpacks share a branch. Loop-return test gains an
elemental two-output call after the returning loop.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/jax/tree.py | 29 ++++++++++++++++++-----------
 tests/test_jax_transform.py      | 15 +++++++++++++--
 2 files changed, 31 insertions(+), 13 deletions(-)

diff --git a/src/recast/transform/jax/tree.py b/src/recast/transform/jax/tree.py
index 5c48199..09c8826 100644
--- a/src/recast/transform/jax/tree.py
+++ b/src/recast/transform/jax/tree.py
@@ -1708,13 +1708,17 @@ def _fold_returns(stmts: list[ast.stmt], rest: list[ast.stmt]) -> list[ast.stmt]
 
 def _guard_after_returns(stmts: list[ast.stmt]) -> list[ast.stmt]:
     """``return`` -> ``_ret = True``; what follows a statement that may have
-    returned, in the same block, runs under ``if not _ret`` -- recursively
-    through branches and loop bodies."""
+    returned, in the same block, runs under one ``if not _ret`` -- the whole
+    remainder as one block, so a call's result buffer (``_out``, never a
+    carry) and its unpacks stay together -- recursively through branches
+    and loop bodies."""
     out: list[ast.stmt] = []
-    exited = False
-    for statement in stmts:
+    for at, statement in enumerate(stmts):
         had_return = _has_return([statement])  # before the rewrite takes the returns away
         if isinstance(statement, ast.Return):
+            # ``jnp.bool_(True)``, not the literal: a literal store is a
+            # trace-time constant to the backend, which would not carry the
+            # flag out of the loop or the branch that set it.
             rewritten: ast.stmt = ast.copy_location(
                 ast.Assign(
                     targets=[ast.Name(id="_ret", ctx=ast.Store())],
@@ -1728,15 +1732,18 @@ def _guard_after_returns(stmts: list[ast.stmt]) -> list[ast.stmt]:
                 inner = getattr(statement, field, None)
                 if isinstance(inner, list) and inner:
                     setattr(statement, field, _guard_after_returns(inner))
-        if exited:
-            rewritten = ast.If(
-                test=ast.UnaryOp(op=ast.Not(), operand=ast.Name(id="_ret", ctx=ast.Load())),
-                body=[rewritten],
-                orelse=[],
-            )
         out.append(rewritten)
         if had_return:
-            exited = True
+            rest = _guard_after_returns(stmts[at + 1 :])
+            if rest:
+                out.append(
+                    ast.If(
+                        test=ast.UnaryOp(op=ast.Not(), operand=ast.Name(id="_ret", ctx=ast.Load())),
+                        body=rest,
+                        orelse=[],
+                    )
+                )
+            return out
     return out
 
 
diff --git a/tests/test_jax_transform.py b/tests/test_jax_transform.py
index 42bb095..c39df16 100644
--- a/tests/test_jax_transform.py
+++ b/tests/test_jax_transform.py
@@ -822,6 +822,7 @@ def run(bad, worse):
     real(8), intent(out), dimension(ngrdcol, nzt) :: y
     integer, intent(in) :: err(ngrdcol)
     integer :: i
+    real(8) :: lo, hi
     x = x * c%coef(:, 1:nzt)
     y = x
     do i = 1, ngrdcol
@@ -831,7 +832,15 @@ def run(bad, worse):
       y(i, :) = y(i, :) + 1.0d0
     end do
     y = y + 10.0d0
+    call bracket( x(1, 1), lo, hi )
+    y = y + lo + hi
   end subroutine apply
+  elemental subroutine bracket( v, lo, hi )
+    real(8), intent(in)  :: v
+    real(8), intent(out) :: lo, hi
+    lo = v - 1.0d0
+    hi = v + 1.0d0
+  end subroutine bracket
 end module loopret_mod
 """
 
@@ -870,8 +879,10 @@ def run(err):
             )
             return np.asarray(y).ravel().tolist()
 
-        assert run([0, 0]) == [13.0, 13.0]
-        assert run([0, 1]) == [3.0, 2.0]  # the second column returns before its +1 and the +10
+        # x(1,1) is 2 after scaling: lo + hi = 4 on the path that reaches the
+        # bracket call (an elemental with two outputs, the anchor's _out).
+        assert run([0, 0]) == [17.0, 17.0]
+        assert run([0, 1]) == [3.0, 2.0]  # the second column returns before its +1 and the rest
         assert run([1, 0]) == [2.0, 2.0]
     finally:
         sys.path.remove(str(out))

From d7ded891012cac36803d8c494a53022446cf3ef8 Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Sat, 5 Sep 2026 01:08:24 -0600
Subject: [PATCH 64/73] jax: rebuild a dummy object handed whole to the host at
 kernel entry

A kernel takes a dummy object apart into components, but its body may
still hand the object whole to something on the host: a framework
stand-in's query on it, a host-kept check under its guard. The name was
never bound in the kernel (UnboundLocalError on the first such use, in
the whole-step port). Rebuild it once at entry from the components, the
way the NumPy flat wrapper does before its call; absent optionals stay
None so present() stays false.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/jax/tree.py | 50 ++++++++++++++++++++-
 tests/test_jax_transform.py      | 77 ++++++++++++++++++++++++++++++++
 2 files changed, 126 insertions(+), 1 deletion(-)

diff --git a/src/recast/transform/jax/tree.py b/src/recast/transform/jax/tree.py
index 09c8826..4536c1e 100644
--- a/src/recast/transform/jax/tree.py
+++ b/src/recast/transform/jax/tree.py
@@ -1646,7 +1646,55 @@ def _rewritten_body(fn: ast.FunctionDef, rewrite: _Rewrite) -> list[ast.stmt]:
                 and rewrite._is_dynamic(node.upper)
             ):
                 raise NotFlat(f"a dynamic slice outside a store or a sum: {ast.unparse(node)}")
-    return lowered
+    return [*_rebuilt_objects(lowered, rewrite), *lowered]
+
+
+def _rebuilt_objects(lowered: list[ast.stmt], rewrite: _Rewrite) -> list[ast.stmt]:
+    """A dummy object the body still hands whole to the host -- a stand-in's
+    query on it (``var_on_stats_list(stats, name)``), a check the port left
+    on the host under its guard -- is rebuilt once, at entry, from the
+    components the kernel takes, the way the NumPy flat wrapper rebuilds
+    it before its call. A Python record at trace time; the components are
+    what flow, and nothing carries the record."""
+    if rewrite.plan is None:
+        return []
+    # An optional dummy the plan leaves out is absent (``k = None`` at the
+    # top, by _absent_optionals): the body's ``present(k)`` must stay false.
+    optional = {
+        _py(a["name"]) for a in rewrite.plan.subprogram.get("args") or () if a.get("optional")
+    }
+    whole = {
+        node.id
+        for statement in lowered
+        for node in ast.walk(statement)
+        if isinstance(node, ast.Name) and isinstance(node.ctx, ast.Load)
+    }
+    rebuilt: list[ast.stmt] = []
+    for obj in rewrite.plan.objects:
+        if obj.kind != "dummy" or obj.name not in whole or obj.name in optional:
+            continue
+        if not obj.components:
+            continue
+        rebuilt.append(
+            ast.fix_missing_locations(
+                ast.Assign(
+                    targets=[ast.Name(id=obj.name, ctx=ast.Store())],
+                    value=ast.Call(
+                        func=ast.Attribute(
+                            value=ast.Name(id="_host", ctx=ast.Load()),
+                            attr="_Record",
+                            ctx=ast.Load(),
+                        ),
+                        args=[],
+                        keywords=[
+                            ast.keyword(arg=c.name, value=ast.Name(id=c.flat, ctx=ast.Load()))
+                            for c in obj.components
+                        ],
+                    ),
+                )
+            )
+        )
+    return rebuilt
 
 
 def _dim_sources(args: list[dict[str, Any]]) -> dict[str, tuple[str, int]]:
diff --git a/tests/test_jax_transform.py b/tests/test_jax_transform.py
index c39df16..b75e52c 100644
--- a/tests/test_jax_transform.py
+++ b/tests/test_jax_transform.py
@@ -1113,6 +1113,83 @@ def test_a_flat_companion_the_port_left_on_the_host_stays_under_its_guard(tmp_pa
                 sys.modules.pop(name, None)
 
 
+OBJECT_QUERY = """\
+module oquery_mod
+  use ocheck_mod, only: knobs_type
+  use stats_query_mod, only: var_on_list
+  implicit none
+contains
+  subroutine step( n, k, x, y )
+    integer, intent(in) :: n
+    type(knobs_type), intent(in) :: k
+    real(8), intent(in) :: x(n)
+    real(8), intent(out) :: y(n)
+    y = k%gain * x
+    if ( var_on_list( k, "gain" ) ) then
+      y = y + k%gain
+    end if
+  end subroutine step
+end module oquery_mod
+"""
+
+STATS_QUERY = """\
+module stats_query_mod
+  use ocheck_mod, only: knobs_type
+  implicit none
+contains
+  logical function var_on_list( k, name )
+    type(knobs_type), intent(in) :: k
+    character(len=*), intent(in) :: name
+    var_on_list = .true.
+  end function var_on_list
+end module stats_query_mod
+"""
+
+
+def test_an_object_handed_whole_to_the_host_is_rebuilt_at_entry(tmp_path: Path) -> None:
+    """``if ( var_on_stats_list( stats, "rsat" ) )``: the query is a
+    framework stand-in's, and it takes the object whole -- which the kernel
+    took apart into components. Rebuilt once at entry from them, as the
+    NumPy flat wrapper does, instead of an UnboundLocalError on a name the
+    kernel never bound."""
+    import importlib
+    import sys
+
+    from recast.transform.jax.tree import TreeToJax
+    from recast.transform.numpy.tree import TreeConventions
+
+    (tmp_path / "ocheck_mod.f90").write_text(OBJECT_CHECK)
+    (tmp_path / "oquery_mod.f90").write_text(OBJECT_QUERY)
+    (tmp_path / "stats_query_mod.f90").write_text(STATS_QUERY)
+    frontend = FortranFrontend(flatten=True, stub_modules=["stats_query_mod"])
+    unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:oquery_mod")
+    facts = frontend.analyze(unit, tmp_path)
+    conventions = TreeConventions(
+        stub_modules=frozenset({"stats_query_mod"}),
+        framework={"stats_query_mod": "def var_on_list(k, name):\n    return name == 'gain'\n"},
+    )
+    candidate = TreeToJax(conventions).apply(unit, facts, {"root": str(tmp_path)})
+    assert "step_flat" in candidate.notes["jax"]["kernels"], candidate.notes["jax"]
+    emitted = candidate.files[Path("oquery_mod_jax.py")].decode()
+    assert "k = _host._Record(gain=k__gain)" in emitted
+    out = tmp_path / "emitted"
+    out.mkdir()
+    for path, content in candidate.files.items():
+        (out / path.name).write_bytes(content)
+    sys.path.insert(0, str(out))
+    try:
+        module = importlib.import_module("oquery_mod_jax")
+        import numpy as np
+
+        got = module.step_flat(2, np.array([1.0, -1.0]), np.zeros(2), 2, np.float64(2.0))
+        assert np.asarray(got).tolist() == [4.0, 0.0]
+    finally:
+        sys.path.remove(str(out))
+        for name in list(sys.modules):
+            if name.startswith(("oquery_mod", "ocheck_mod", "stats_query_mod")):
+                sys.modules.pop(name, None)
+
+
 SILENT_CHECK = """\
 module silent_mod
   implicit none

From 18c6303232c8b899d6c902d7ccf6bf9e2a1daa38 Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Sat, 5 Sep 2026 01:42:09 -0600
Subject: [PATCH 65/73] jax: six rules the whole-step port needed

- a loop whose body lowered to nothing (statistics sampling dropped by
  its stand-in) is dropped, not refused for carrying nothing
- a logical scalar dummy is static under jit: a configuration switch,
  whose dead branch is never traced
- a named or computed stride (a grid's direction, +-1 at run time):
  trips counted by the runtime's _f_trips, the index remapped
- an output the anchor discards (_ = _out[k]) is dropped from the
  elided buffer instead of refused for having no actual
- an optional object the caller leaves out is absent in the callee:
  None in its component slots, its outputs dropped; a call that leaves
  out every (optional) object still reaches the flat function
- a callee's extent argument (an allocatable component sized by an
  allocating routine's dummy) is the caller's own extent argument for
  that axis, else that axis of the caller's component

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/jax/backend.py |  65 ++++-
 src/recast/transform/jax/runtime.py |  11 +
 src/recast/transform/jax/tree.py    |  64 ++++-
 tests/test_jax_transform.py         | 359 ++++++++++++++++++++++++++++
 4 files changed, 494 insertions(+), 5 deletions(-)

diff --git a/src/recast/transform/jax/backend.py b/src/recast/transform/jax/backend.py
index 8c8049f..c2244b2 100644
--- a/src/recast/transform/jax/backend.py
+++ b/src/recast/transform/jax/backend.py
@@ -659,26 +659,39 @@ def lower_for(self, s, depth):
             and not it.keywords
         ):
             raise JaxQueue("non-range for")
-        step = 1
+        step: int | None = 1
         # Annotated because the first branch would otherwise pin these to
         # Constant and the others assign general expressions to them.
         lo: ast.expr
         hi: ast.expr
+        stride: ast.expr | None = None
         if len(it.args) == 1:
             lo, hi = ast.Constant(value=0), it.args[0]
         elif len(it.args) == 2:
             lo, hi = it.args
         elif len(it.args) == 3:
-            step = _const_int(it.args[2])
             lo, hi = it.args[0], it.args[1]
+            try:
+                step = _const_int(it.args[2])
+            except JaxQueue:
+                step = None
             if step not in (1, -1):
-                raise JaxQueue("range step not +-1")
+                # ``do k = lb, ub, dir``: a stride that is a name or an
+                # expression (a grid's direction, +-1 at run time). Trips
+                # counted by the runtime, the index remapped from the
+                # trip: k = lo + t * stride.
+                step, stride = None, it.args[2]
         else:
             raise JaxQueue("malformed range")
 
         bound_before = set(self.bound)
         self.bound.add(s.target.id)
         body = self.lower_block(_cycle_to_else(s.body, []), depth + 1)
+        if not body or all(isinstance(b, ast.Pass) for b in body):
+            # A loop whose body lowered to nothing: statistics calls the
+            # stand-in dropped, a nested loop that went the same way.
+            # Fortran ran it for nothing; nothing is what it becomes.
+            return []
         settled = _trace_constant_stores(body)
         carried = [
             n
@@ -704,6 +717,45 @@ def lower_for(self, s, depth):
             carry_call.args = [lo, hi, ast.Name(id=fname, ctx=ast.Load()), init]
             return [fn, result]
 
+        if stride is not None:
+            lo_name, st_name, cnt_name = f"_lo_{self.n}", f"_st_{self.n}", f"_cnt_{self.n}"
+            pre = [
+                ast.Assign(targets=[ast.Name(id=lo_name, ctx=ast.Store())], value=lo),
+                ast.Assign(targets=[ast.Name(id=st_name, ctx=ast.Store())], value=stride),
+                ast.Assign(
+                    targets=[ast.Name(id=cnt_name, ctx=ast.Store())],
+                    value=ast.Call(
+                        func=ast.Name(id="_f_trips", ctx=ast.Load()),
+                        args=[
+                            ast.Name(id=lo_name, ctx=ast.Load()),
+                            hi,
+                            ast.Name(id=st_name, ctx=ast.Load()),
+                        ],
+                        keywords=[],
+                    ),
+                ),
+            ]
+            remap = ast.Assign(
+                targets=[ast.Name(id=s.target.id, ctx=ast.Store())],
+                value=ast.BinOp(
+                    left=ast.Name(id=lo_name, ctx=ast.Load()),
+                    op=ast.Add(),
+                    right=ast.BinOp(
+                        left=ast.Name(id="_r", ctx=ast.Load()),
+                        op=ast.Mult(),
+                        right=ast.Name(id=st_name, ctx=ast.Load()),
+                    ),
+                ),
+            )
+            fn = self._carry_fn(fname, ["_r", "_c"], carried, [remap, *body])
+            carry_call.args = [
+                ast.Constant(value=0),
+                ast.Name(id=cnt_name, ctx=ast.Load()),
+                ast.Name(id=fname, ctx=ast.Load()),
+                init,
+            ]
+            return [*pre, fn, result]
+
         # step -1: Fortran DO k=hi,lo,-1 arrived as range(hi, stop, -1)
         # iterating hi..stop+1. Remap: t in [0, hi-stop), k = hi - t.
         # Bounds are hoisted (evaluated once, Fortran DO semantics).
@@ -991,7 +1043,12 @@ def static_spec(fn_src, sub, traced_scalars=frozenset()):
             # when it is module state (the ``__`` spelling: dtime_ml, dtstep
             # -- namelist configuration): an ordinary real dummy (xa, xb,
             # tol) must stay traced, or jvp/grad against it breaks.
-            if r["dtype"] in ("int32", "str") or (
+            # A logical scalar dummy is a switch (a model's configuration
+            # flags): static, so its branches are Python ifs at trace time
+            # and a path the run never takes is never traced. When another
+            # kernel passes it a tracer, the runtime's ``_f_concrete``
+            # takes the lax.cond form instead.
+            if r["dtype"] in ("int32", "str", "bool") or (
                 r["dtype"] in ("int64", "float64") and "__" in r["name"]
             ):
                 nums.append(pos)
diff --git a/src/recast/transform/jax/runtime.py b/src/recast/transform/jax/runtime.py
index c28ee95..f14dfb6 100644
--- a/src/recast/transform/jax/runtime.py
+++ b/src/recast/transform/jax/runtime.py
@@ -55,6 +55,7 @@
     "_f_sqrt",
     "_f_tiny",
     "_f_trim",
+    "_f_trips",
     "_f_vceil",
     "_f_vdot",
     "_f_verf",
@@ -178,6 +179,16 @@ def _f_fori(lo, hi, body, init):
     return lax.fori_loop(lo, hi, body, init)
 
 
+def _f_trips(lo, hi, step):
+    """How many times ``range(lo, hi, step)`` runs -- Fortran's DO trip
+    count, ``max(0, (hi - lo + step - sign(step)) // step)`` -- as a
+    Python int when every bound is, else traced."""
+    static = (int, np.integer)
+    if isinstance(lo, static) and isinstance(hi, static) and isinstance(step, static):
+        return len(range(int(lo), int(hi), int(step)))
+    return jnp.maximum(0, (hi - lo + step - jnp.sign(step)) // step)
+
+
 def _f_ecall(fn, *args, **kw):
     """ELEMENTAL procedure broadcast over array actuals: the scalar kernel
     per element, in sequence, as the NumPy runtime's np.vectorize runs it.
diff --git a/src/recast/transform/jax/tree.py b/src/recast/transform/jax/tree.py
index 4536c1e..6b50c04 100644
--- a/src/recast/transform/jax/tree.py
+++ b/src/recast/transform/jax/tree.py
@@ -374,6 +374,10 @@ def visit_Assign(self, node: ast.Assign) -> Any:
             if at >= len(actuals):
                 raise NotFlat(f"{ast.unparse(node)}: the elided buffer has no slot {at}")
             target = node.targets[0]
+            if isinstance(target, ast.Name) and target.id == "_":
+                # ``_ = _out[4]``: an output the anchor discards (an optional
+                # OUT the caller did not ask for, absent on both sides).
+                return None
             if isinstance(target, ast.Name) and self.spelling.object_of(target) is not None:
                 return None  # an object: its components bound at the call
             actual = actuals[at]
@@ -961,6 +965,13 @@ def _flat_callee(self, call: ast.Call) -> FlatPlan | None:
             self.spelling.object_of(k.value) is not None for k in call.keywords
         ):
             return callee
+        optional = {a["name"].lower() for a in callee.subprogram["args"] if a.get("optional")}
+        dummies = [obj for obj in callee.objects if obj.kind == "dummy"]
+        if dummies and all(obj.name in optional for obj in dummies):
+            # Every object the callee takes is optional and this call
+            # leaves them all out (pdf_closure without its implicit
+            # coefficients): still the flat function, with them absent.
+            return callee
         return None
 
     def _specialize(self, name: str, call: ast.Call, source: dict[str, Any]) -> FlatPlan | None:
@@ -1226,10 +1237,21 @@ def _rewrite_call(self, target: ast.expr | None, call: ast.Call) -> Any:
                 actual_by_dummy[keyword_.arg.lower().rstrip("_")] = keyword_.value
         # Which callee object is which caller object.
         objects: dict[str, str] = {}
+        optional_dummies = {
+            a["name"].lower() for a in callee.subprogram["args"] if a.get("optional")
+        }
+        # An optional object the caller leaves out: the callee's kernel
+        # still takes its components (the flat signature has them), and
+        # its ``present(obj)`` is false at trace time. ``None`` in every
+        # slot, and whatever the kernel hands back for them is dropped.
+        absent: set[str] = set()
         for obj in callee.objects:
             if obj.kind == "dummy":
                 actual = actual_by_dummy.get(obj.name)
                 passed: str | None = self.spelling.object_of(actual) if actual is not None else None
+                if passed is None and actual is None and obj.name in optional_dummies:
+                    absent.add(obj.name)
+                    continue
                 if passed is None:
                     raise NotFlat(f"{callee.subprogram['name']}: object {obj.name} not passed")
                 objects[obj.name] = passed
@@ -1245,6 +1267,44 @@ def _rewrite_call(self, target: ast.expr | None, call: ast.Call) -> Any:
                 continue  # returned, not taken -- the emitted convention
             if name == callee.patch_count:
                 args.append(ast.Name(id=self.plan.patch_count, ctx=ast.Load()))
+            elif "__" in name and name.split("__", 1)[0] in absent:
+                args.append(ast.Constant(value=None))
+            elif name in callee.extent_args:
+                # An extent the callee's plan could not spell (an allocatable
+                # component sized by an allocating routine's dummy): the
+                # caller's own extent argument for the same axis when it
+                # has one, else that axis of the caller's component, a
+                # static shape at trace time.
+                owner, comp, axis = callee.extent_args[name]
+                passed_obj = objects.get(owner)
+                own = next(
+                    (
+                        n
+                        for n, (o, c, ax) in self.plan.extent_args.items()
+                        if o == passed_obj and c == comp and ax == axis
+                    ),
+                    None,
+                )
+                if own is not None:
+                    args.append(ast.Name(id=own, ctx=ast.Load()))
+                    continue
+                sized = caller_components.get(passed_obj or "", {}).get(comp)
+                if sized is None:
+                    raise NotFlat(
+                        f"{callee.subprogram['name']}: extent {name} of {owner}%{comp}"
+                        " not in the caller's plan"
+                    )
+                args.append(
+                    ast.Subscript(
+                        value=ast.Attribute(
+                            value=ast.Name(id=sized, ctx=ast.Load()),
+                            attr="shape",
+                            ctx=ast.Load(),
+                        ),
+                        slice=ast.Constant(value=int(axis) - 1),  # the plan's axis is 1-based
+                        ctx=ast.Load(),
+                    )
+                )
             elif "__" in name and (owner := name.split("__", 1)[0]) in objects:
                 comp = name.split("__", 1)[1]
                 flat_name: str | None = caller_components.get(objects[owner], {}).get(comp)
@@ -1342,7 +1402,9 @@ def _rewrite_call(self, target: ast.expr | None, call: ast.Call) -> Any:
         targets: list[ast.expr] = []
         follow: list[ast.stmt] = []
         for name in _outputs(callee):
-            if "__" in name and (owner := name.split("__", 1)[0]) in objects:
+            if "__" in name and name.split("__", 1)[0] in absent:
+                targets.append(ast.Name(id="_", ctx=ast.Store()))
+            elif "__" in name and (owner := name.split("__", 1)[0]) in objects:
                 comp = name.split("__", 1)[1]
                 targets.append(
                     ast.Name(id=caller_components[objects[owner]][comp], ctx=ast.Store())
diff --git a/tests/test_jax_transform.py b/tests/test_jax_transform.py
index b75e52c..07d1cfe 100644
--- a/tests/test_jax_transform.py
+++ b/tests/test_jax_transform.py
@@ -1190,6 +1190,365 @@ def test_an_object_handed_whole_to_the_host_is_rebuilt_at_entry(tmp_path: Path)
                 sys.modules.pop(name, None)
 
 
+EMPTY_LOOP = """\
+module emptyloop_mod
+  implicit none
+contains
+  subroutine scale( n, x, y )
+    integer, intent(in) :: n
+    real(8), intent(in) :: x(n)
+    real(8), intent(out) :: y(n)
+    integer :: i
+    y = 2.0d0 * x
+    do i = 1, n
+      continue
+    end do
+  end subroutine scale
+end module emptyloop_mod
+"""
+
+
+def test_a_loop_whose_body_lowered_to_nothing_is_dropped(tmp_path: Path) -> None:
+    """CLUBB's clipping routines loop over the columns to sample statistics;
+    with the sampling calls dropped by their stand-in, the loop's body is
+    ``pass`` and there is nothing to carry. Fortran ran it for nothing."""
+    import importlib
+    import sys
+
+    candidate = port(tmp_path, EMPTY_LOOP, "emptyloop_mod")
+    assert candidate.notes["jax"]["kernels"] == ["scale"], candidate.notes["jax"]
+    out = tmp_path / "emitted"
+    out.mkdir()
+    for path, content in candidate.files.items():
+        (out / path.name).write_bytes(content)
+    sys.path.insert(0, str(out))
+    try:
+        module = importlib.import_module("emptyloop_mod_jax")
+        import numpy as np
+
+        got = module.scale(2, np.array([1.0, 3.0]))
+        assert np.asarray(got).tolist() == [2.0, 6.0]
+    finally:
+        sys.path.remove(str(out))
+        for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"):
+            sys.modules.pop(f"emptyloop_mod{suffix}", None)
+
+
+STRIDED = """\
+module strided_mod
+  implicit none
+contains
+  subroutine running( n, dir, x, y )
+    integer, intent(in) :: n, dir
+    real(8), intent(in) :: x(n)
+    real(8), intent(out) :: y(n)
+    integer :: k, lb, ub
+    real(8) :: acc
+    if ( dir > 0 ) then
+      lb = 1
+      ub = n
+    else
+      lb = n
+      ub = 1
+    end if
+    acc = 0.0d0
+    do k = lb, ub, dir
+      acc = acc + x(k)
+      y(k) = acc
+    end do
+  end subroutine running
+end module strided_mod
+"""
+
+
+def test_a_loop_with_a_named_stride_runs_in_that_direction(tmp_path: Path) -> None:
+    """``do k = gr%k_lb_zt, gr%k_ub_zt, gr%grid_dir_indx``: CLUBB's grid
+    direction is a run-time +-1, a name in the stride slot. The trip count
+    is the runtime's, the index is remapped from the trip."""
+    import importlib
+    import sys
+
+    candidate = port(tmp_path, STRIDED, "strided_mod")
+    assert candidate.notes["jax"]["kernels"] == ["running"], candidate.notes["jax"]
+    assert "_f_trips(" in candidate.files[Path("strided_mod_jax.py")].decode()
+    out = tmp_path / "emitted"
+    out.mkdir()
+    for path, content in candidate.files.items():
+        (out / path.name).write_bytes(content)
+    sys.path.insert(0, str(out))
+    try:
+        module = importlib.import_module("strided_mod_jax")
+        import numpy as np
+
+        x = np.array([1.0, 2.0, 4.0])
+        up = np.asarray(module.running(3, 1, x)).tolist()
+        down = np.asarray(module.running(3, -1, x)).tolist()
+        assert up == [1.0, 3.0, 7.0]
+        assert down == [7.0, 6.0, 4.0]
+    finally:
+        sys.path.remove(str(out))
+        for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"):
+            sys.modules.pop(f"strided_mod{suffix}", None)
+
+
+SWITCHED_CHECK = """\
+module switched_mod
+  use silent_mod, only: complain
+  implicit none
+contains
+  subroutine step( n, l_check, x, y )
+    integer, intent(in) :: n
+    logical, intent(in) :: l_check
+    real(8), intent(in) :: x(n)
+    real(8), intent(out) :: y(n)
+    y = 2.0d0 * x
+    if ( l_check ) then
+      call complain( n, y )
+    end if
+  end subroutine step
+end module switched_mod
+"""
+
+
+def test_a_logical_scalar_dummy_is_a_static_switch(tmp_path: Path) -> None:
+    """CLUBB's configuration flags are logical dummies (the model's
+    ``clubb_config_flags`` components). Static under jit: the branch is a
+    Python if at trace time, and the check the port left on the host --
+    under a flag the run has off -- is never traced."""
+    import importlib
+    import sys
+
+    from recast.transform.jax.tree import TreeToJax
+    from recast.transform.numpy.tree import TreeConventions
+
+    (tmp_path / "silent_mod.f90").write_text(SILENT_CHECK)
+    (tmp_path / "switched_mod.f90").write_text(SWITCHED_CHECK)
+    frontend = FortranFrontend(flatten=True)
+    unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:switched_mod")
+    facts = frontend.analyze(unit, tmp_path)
+    candidate = TreeToJax(TreeConventions()).apply(unit, facts, {"root": str(tmp_path)})
+    assert "step" in candidate.notes["jax"]["kernels"], candidate.notes["jax"]
+    emitted = candidate.files[Path("switched_mod_jax.py")].decode()
+    assert "static_argnums=(0, 1)" in emitted
+    out = tmp_path / "emitted"
+    out.mkdir()
+    for path, content in candidate.files.items():
+        (out / path.name).write_bytes(content)
+    sys.path.insert(0, str(out))
+    try:
+        module = importlib.import_module("switched_mod_jax")
+        import numpy as np
+
+        got = module.step(2, False, np.array([1.0, -1.0]))
+        assert np.asarray(got).tolist() == [2.0, -2.0]
+    finally:
+        sys.path.remove(str(out))
+        for name in list(sys.modules):
+            if name.startswith(("switched_mod", "silent_mod")):
+                sys.modules.pop(name, None)
+
+
+OMITS_AN_OBJECT = """\
+module omit_mod
+  implicit none
+  type knobs_type
+    real(8) :: gain = 2.0d0
+  end type knobs_type
+contains
+  subroutine scale( nzt, ngrdcol, x, y, k )
+    integer, intent(in) :: nzt, ngrdcol
+    real(8), intent(inout), dimension(ngrdcol, nzt) :: x
+    real(8), intent(inout), dimension(ngrdcol, nzt) :: y
+    type(knobs_type), intent(inout), optional :: k
+    if ( present( k ) ) k%gain = k%gain * 2.0d0
+    x = x * 2.0d0
+    y = x + 1.0d0
+  end subroutine scale
+  subroutine step( nzt, ngrdcol, k, x, y )
+    integer, intent(in) :: nzt, ngrdcol
+    type(knobs_type), intent(inout) :: k
+    real(8), intent(inout), dimension(ngrdcol, nzt) :: x
+    real(8), intent(inout), dimension(ngrdcol, nzt) :: y
+    call scale( nzt, ngrdcol, x, y )
+    y = y + k%gain
+  end subroutine step
+end module omit_mod
+"""
+
+
+def test_an_optional_object_the_caller_leaves_out_is_absent_in_the_callee(tmp_path: Path) -> None:
+    """pdf_closure_driver_zm calls pdf_closure without its optional
+    ``pdf_implicit_coefs_terms``. The callee's kernel still takes the
+    object's components (its flat signature has them): ``None`` in each,
+    ``present()`` false at trace time, and what it hands back for them
+    dropped."""
+    import importlib
+    import sys
+
+    from recast.transform.jax.tree import TreeToJax
+    from recast.transform.numpy.tree import TreeConventions
+
+    (tmp_path / "omit_mod.f90").write_text(OMITS_AN_OBJECT)
+    frontend = FortranFrontend(flatten={"patch_count": "ngrdcol", "bounds_pattern": r"^ngrdcol$"})
+    unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:omit_mod")
+    facts = frontend.analyze(unit, tmp_path)
+    candidate = TreeToJax(TreeConventions()).apply(unit, facts, {"root": str(tmp_path)})
+    assert "step_flat" in candidate.notes["jax"]["kernels"], candidate.notes["jax"]
+    out = tmp_path / "emitted"
+    out.mkdir()
+    for path, content in candidate.files.items():
+        (out / path.name).write_bytes(content)
+    sys.path.insert(0, str(out))
+    try:
+        module = importlib.import_module("omit_mod_jax")
+        import numpy as np
+
+        x = np.ones((1, 2), order="F")
+        result = module.step_flat(2, 1, x, np.zeros((1, 2), order="F"), np.float64(2.0))
+        got = [np.asarray(v).tolist() for v in result]
+        # gain is read, never written, so it is not returned; the callee's
+        # present(k) branch did not run: y is x * 2 + 1 + 2.
+        assert got == [[[2.0, 2.0]], [[5.0, 5.0]]]
+    finally:
+        sys.path.remove(str(out))
+        for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"):
+            sys.modules.pop(f"omit_mod{suffix}", None)
+
+
+DISCARDS_AN_OUTPUT = """\
+module discard_mod
+  implicit none
+  type knobs_type
+    real(8) :: gain = 2.0d0
+  end type knobs_type
+contains
+  subroutine solve( nzt, ngrdcol, k, x, y, resid )
+    integer, intent(in) :: nzt, ngrdcol
+    type(knobs_type), intent(in) :: k
+    real(8), intent(in), dimension(ngrdcol, nzt) :: x
+    real(8), intent(inout), dimension(ngrdcol, nzt) :: y
+    real(8), intent(out), dimension(ngrdcol, nzt), optional :: resid
+    y = k%gain * x
+    if ( present( resid ) ) resid = y - x
+  end subroutine solve
+  subroutine step( nzt, ngrdcol, k, x, y )
+    integer, intent(in) :: nzt, ngrdcol
+    type(knobs_type), intent(in) :: k
+    real(8), intent(in), dimension(ngrdcol, nzt) :: x
+    real(8), intent(inout), dimension(ngrdcol, nzt) :: y
+    call solve( nzt, ngrdcol, k, x, y )
+    y = y + 1.0d0
+  end subroutine step
+end module discard_mod
+"""
+
+
+def test_an_output_the_anchor_discards_is_dropped_from_the_elided_buffer(tmp_path: Path) -> None:
+    """CLUBB's solvers return an optional residual the caller does not ask
+    for: the anchor's ``_ = _out[4]`` names a slot with no actual, and the
+    unpack is dropped rather than refused."""
+    import importlib
+    import sys
+
+    from recast.transform.jax.tree import TreeToJax
+    from recast.transform.numpy.tree import TreeConventions
+
+    (tmp_path / "discard_mod.f90").write_text(DISCARDS_AN_OUTPUT)
+    frontend = FortranFrontend(flatten={"patch_count": "ngrdcol", "bounds_pattern": r"^ngrdcol$"})
+    unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:discard_mod")
+    facts = frontend.analyze(unit, tmp_path)
+    candidate = TreeToJax(TreeConventions()).apply(unit, facts, {"root": str(tmp_path)})
+    assert "_ = _out[" in candidate.files[Path("discard_mod_numpy.py")].decode()
+    assert "step_flat" in candidate.notes["jax"]["kernels"], candidate.notes["jax"]
+    out = tmp_path / "emitted"
+    out.mkdir()
+    for path, content in candidate.files.items():
+        (out / path.name).write_bytes(content)
+    sys.path.insert(0, str(out))
+    try:
+        module = importlib.import_module("discard_mod_jax")
+        import numpy as np
+
+        x = np.ones((1, 2), order="F")
+        got = module.step_flat(2, 1, x, np.zeros((1, 2), order="F"), np.float64(3.0))
+        assert np.asarray(got).tolist() == [[4.0, 4.0]]
+    finally:
+        sys.path.remove(str(out))
+        for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"):
+            sys.modules.pop(f"discard_mod{suffix}", None)
+
+
+SIZED_BY_AN_ALLOCATOR = """\
+module sized_mod
+  implicit none
+  type coefs_type
+    real(8), allocatable :: coef(:, :)
+  end type coefs_type
+contains
+  subroutine init_coefs( ngrdcol, nz, p )
+    integer, intent(in) :: ngrdcol, nz
+    type(coefs_type), intent(out) :: p
+    allocate( p%coef(1:ngrdcol, 1:nz) )
+    p%coef = 0.0d0
+  end subroutine init_coefs
+  subroutine apply_coefs( ngrdcol, p, x )
+    integer, intent(in) :: ngrdcol
+    type(coefs_type), intent(in) :: p
+    real(8), intent(inout) :: x(ngrdcol)
+    integer :: i
+    do i = 1, ngrdcol
+      x(i) = x(i) + sum( p%coef(i, :) )
+    end do
+  end subroutine apply_coefs
+  subroutine step( ngrdcol, nz, p, x )
+    integer, intent(in) :: ngrdcol, nz
+    type(coefs_type), intent(inout) :: p
+    real(8), intent(inout) :: x(ngrdcol)
+    p%coef(:, 1) = p%coef(:, 1) + 1.0d0
+    call apply_coefs( ngrdcol, p, x )
+  end subroutine step
+end module sized_mod
+"""
+
+
+def test_a_callee_extent_argument_is_the_callers_axis(tmp_path: Path) -> None:
+    """A component allocated by another routine's dummy (CLUBB's
+    ``coef_wp4_implicit(1:ngrdcol, 1:nz)``) reaches a callee that has no
+    dummy of that name as an extent argument of its plan. The caller
+    passes that axis of its own component."""
+    import importlib
+    import sys
+
+    from recast.fortran.flatten import plans_from_facts
+    from recast.transform.jax.tree import TreeToJax
+    from recast.transform.numpy.tree import TreeConventions
+
+    (tmp_path / "sized_mod.f90").write_text(SIZED_BY_AN_ALLOCATOR)
+    frontend = FortranFrontend(flatten={"patch_count": "ngrdcol", "bounds_pattern": r"^ngrdcol$"})
+    unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:sized_mod")
+    facts = frontend.analyze(unit, tmp_path)
+    plans = {p.name: p for p in plans_from_facts(facts)}
+    assert plans["apply_coefs_flat"].extent_args, plans["apply_coefs_flat"]
+    candidate = TreeToJax(TreeConventions()).apply(unit, facts, {"root": str(tmp_path)})
+    assert "step_flat" in candidate.notes["jax"]["kernels"], candidate.notes["jax"]
+    out = tmp_path / "emitted"
+    out.mkdir()
+    for path, content in candidate.files.items():
+        (out / path.name).write_bytes(content)
+    sys.path.insert(0, str(out))
+    try:
+        module = importlib.import_module("sized_mod_jax")
+        import numpy as np
+
+        coef = np.zeros((1, 2), order="F")
+        got = module.step_flat(1, 2, np.array([1.0]), coef)
+        assert [np.asarray(v).tolist() for v in got] == [[2.0], [[1.0, 0.0]]]
+    finally:
+        sys.path.remove(str(out))
+        for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"):
+            sys.modules.pop(f"sized_mod{suffix}", None)
+
+
 SILENT_CHECK = """\
 module silent_mod
   implicit none

From d4b732ed40a5011ceb2b144c14ff91b7506bef04 Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Sat, 5 Sep 2026 01:56:09 -0600
Subject: [PATCH 66/73] jax: EXIT as a carried flag, for-else, own subprograms
 kept on the host

- a break inside a lowered loop is a flag the loop carries: set where
  the break was, it guards the rest of that trip and every later one;
  the DO variable's value at the exit is kept (int32) and rebound after
  the loop; a for-else (the DO ran to completion) runs under the flag
  inverted, its completion value cast to the index's dtype
- static loops scan over int32 indices (Fortran's default integer)
  instead of fori_loop's default-int scan, so a store of the index into
  an integer local or dummy keeps its dtype across a lax.cond
- a subprogram of the module the port could not emit, called from one
  it can, stays on the host (_host.<name>) under the anchor's guard and
  is named in the notes' host_calls, instead of delegating its caller
  and the callers above it; build_module returns those calls
- _lo_n/_st_n hoisted stride bounds are never carries

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/jax/backend.py   | 255 ++++++++++++++++++++++++--
 src/recast/transform/jax/runtime.py   |  18 +-
 src/recast/transform/jax/translate.py |   5 +-
 src/recast/transform/jax/tree.py      |   7 +-
 tests/test_jax_transform.py           |  90 ++++++++-
 5 files changed, 352 insertions(+), 23 deletions(-)

diff --git a/src/recast/transform/jax/backend.py b/src/recast/transform/jax/backend.py
index c2244b2..14847d3 100644
--- a/src/recast/transform/jax/backend.py
+++ b/src/recast/transform/jax/backend.py
@@ -267,6 +267,7 @@ class CallRewrite(ast.NodeTransformer):
     def __init__(self, call_map, known_subs):
         self.map = call_map
         self.subs = known_subs
+        self.host_calls: list[str] = []
 
     def visit_Call(self, node):
         self.generic_visit(node)
@@ -300,7 +301,20 @@ def visit_Call(self, node):
                 node.func = ast.Name(id=f"_{f.id}_k_impl", ctx=ast.Load())
                 node.args = pos + [ast.Name(id=c, ctx=ast.Load()) for c in info["closure"]]
             elif f.id in self.subs:
-                raise JaxQueue(f"calls non-emitted subprogram {f.id}")
+                # A subprogram of this module the port could not emit,
+                # called from one it can: the call stays the host's
+                # (``_host.<name>``, the NumPy anchor's), under whatever
+                # guard the anchor put it -- a configuration switch, static
+                # at trace time, is the usual one -- and is never traced
+                # while the guard holds. If the guard ever lets it
+                # through, the trace fails on it by name, which the gate
+                # reports. Delegating the caller, and its callers up to
+                # the driver, for a path the run never takes was the
+                # alternative. The note names it.
+                self.host_calls.append(f.id)
+                node.func = ast.Attribute(
+                    value=ast.Name(id="_host", ctx=ast.Load()), attr=f.id, ctx=ast.Load()
+                )
         return node
 
 
@@ -343,6 +357,109 @@ def _cycle_to_else(stmts, rest):
     return [*out, *rest]
 
 
+def _not_flag(flag: str) -> ast.expr:
+    """``jnp.logical_not(flag)``: a traced test. Python's ``not`` on a
+    tracer is the boolean conversion JAX refuses, and the expression
+    mapping that would spell it ran before the loop was lowered."""
+    return ast.Call(
+        func=ast.Attribute(
+            value=ast.Name(id="jnp", ctx=ast.Load()), attr="logical_not", ctx=ast.Load()
+        ),
+        args=[ast.Name(id=flag, ctx=ast.Load())],
+        keywords=[],
+    )
+
+
+def _jnp_bool(value: bool) -> ast.expr:
+    """``jnp.bool_(value)``: a traced store, not a literal the backend
+    would treat as a trace-time constant and fail to carry."""
+    return ast.Call(
+        func=ast.Attribute(value=ast.Name(id="jnp", ctx=ast.Load()), attr="bool_", ctx=ast.Load()),
+        args=[ast.Constant(value=value)],
+        keywords=[],
+    )
+
+
+def _breaks_at_level(stmts) -> bool:
+    """A ``break`` that belongs to this loop: not one inside a nested loop."""
+    for st in stmts:
+        if isinstance(st, ast.Break):
+            return True
+        if isinstance(st, ast.For | ast.While):
+            continue
+        for field in ("body", "orelse"):
+            inner = getattr(st, field, None)
+            if isinstance(inner, list) and _breaks_at_level(inner):
+                return True
+    return False
+
+
+def _guard_after_breaks(stmts, flag, kept, index):
+    """``break`` -> ``flag = True; kept = index``; what follows a statement
+    that may have broken, in the same block, runs under ``if not flag`` --
+    recursively through branches, never into a nested loop (its breaks are
+    its own)."""
+    out: list[ast.stmt] = []
+    for at, st in enumerate(stmts):
+        had_break = _breaks_at_level([st])
+        if isinstance(st, ast.Break):
+            out.append(
+                ast.copy_location(
+                    ast.Assign(targets=[ast.Name(id=flag, ctx=ast.Store())], value=_jnp_bool(True)),
+                    st,
+                )
+            )
+            rewritten: ast.stmt = ast.copy_location(
+                ast.Assign(
+                    targets=[ast.Name(id=kept, ctx=ast.Store())],
+                    value=ast.Call(
+                        func=ast.Attribute(
+                            value=ast.Call(
+                                func=ast.Attribute(
+                                    value=ast.Name(id="jnp", ctx=ast.Load()),
+                                    attr="asarray",
+                                    ctx=ast.Load(),
+                                ),
+                                args=[ast.Name(id=index, ctx=ast.Load())],
+                                keywords=[],
+                            ),
+                            attr="astype",
+                            ctx=ast.Load(),
+                        ),
+                        args=[
+                            ast.Attribute(
+                                value=ast.Name(id=kept, ctx=ast.Load()),
+                                attr="dtype",
+                                ctx=ast.Load(),
+                            )
+                        ],
+                        keywords=[],
+                    ),
+                ),
+                st,
+            )
+        else:
+            rewritten = st
+            if not isinstance(st, ast.For | ast.While):
+                for field in ("body", "orelse"):
+                    inner = getattr(st, field, None)
+                    if isinstance(inner, list) and inner:
+                        setattr(st, field, _guard_after_breaks(inner, flag, kept, index))
+        out.append(rewritten)
+        if had_break:
+            rest = _guard_after_breaks(stmts[at + 1 :], flag, kept, index)
+            if rest:
+                out.append(
+                    ast.If(
+                        test=_not_flag(flag),
+                        body=rest,
+                        orelse=[],
+                    )
+                )
+            return out
+    return out
+
+
 def _assigned_names(stmts):
     """Names stored by Assign statements, first-assignment order
     (nested fori_loop/cond results arrive as Tuple targets; static
@@ -364,7 +481,9 @@ def add(n):
                 # ``_out`` is the anchor's call-result tuple, assigned and
                 # unpacked on consecutive lines of one block: the block's own,
                 # never a carry (its shape changes from call to call).
-                if isinstance(t, ast.Name) and not re.fullmatch(r"_(?:hi_|cnt_|t)\d+|_out", t.id):
+                if isinstance(t, ast.Name) and not re.fullmatch(
+                    r"_(?:hi_|cnt_|lo_|st_|t)\d+|_out", t.id
+                ):
                     add(t.id)
                 elif isinstance(t, ast.Tuple):
                     for e in t.elts:
@@ -647,8 +766,6 @@ def lower_assign(self, s):
         raise JaxQueue(f"unsupported assign target {type(t).__name__}")
 
     def lower_for(self, s, depth):
-        if s.orelse:
-            raise JaxQueue("for-else")
         if not isinstance(s.target, ast.Name):
             raise JaxQueue("tuple loop target")
         it = s.iter
@@ -684,9 +801,103 @@ def lower_for(self, s, depth):
         else:
             raise JaxQueue("malformed range")
 
+        body_stmts = _cycle_to_else(s.body, [])
+        flag: str | None = None
+        pre_flag: list[ast.stmt] = []
+        if _breaks_at_level(body_stmts):
+            # EXIT: a flag the loop carries. Set where the break was, it
+            # guards the rest of that trip and the whole of every later
+            # one -- the loop runs its trips and does nothing after the
+            # exit, which is what the exit left it doing. A for-else
+            # (the DO ran to completion) runs after the loop under the
+            # same flag, inverted.
+            self.n += 1
+            flag = f"_brk_{self.n}"
+            # The index at the exit is what the DO variable holds after the
+            # loop (CLUBB's window search reads it): captured with the
+            # flag, in the index's own dtype, and rebound after the loop.
+            kept = f"_kx_{self.n}"
+            body_stmts = [
+                ast.If(
+                    test=_not_flag(flag),
+                    body=_guard_after_breaks(body_stmts, flag, kept, s.target.id),
+                    orelse=[],
+                )
+            ]
+            pre_flag = [
+                ast.Assign(targets=[ast.Name(id=flag, ctx=ast.Store())], value=_jnp_bool(False)),
+                ast.Assign(
+                    targets=[ast.Name(id=kept, ctx=ast.Store())],
+                    value=ast.Call(
+                        func=ast.Attribute(
+                            value=ast.Name(id="jnp", ctx=ast.Load()), attr="asarray", ctx=ast.Load()
+                        ),
+                        args=[copy.deepcopy(lo)],
+                        keywords=[
+                            ast.keyword(
+                                arg="dtype",
+                                value=ast.Attribute(
+                                    value=ast.Name(id="jnp", ctx=ast.Load()),
+                                    attr="int32",
+                                    ctx=ast.Load(),
+                                ),
+                            )
+                        ],
+                    ),
+                ),
+            ]
+            self.bound.add(flag)
+            self.bound.add(kept)
+        elif s.orelse:
+            raise JaxQueue("for-else without a break")
         bound_before = set(self.bound)
         self.bound.add(s.target.id)
-        body = self.lower_block(_cycle_to_else(s.body, []), depth + 1)
+        body = self.lower_block(body_stmts, depth + 1)
+        after: list[ast.stmt] = []
+        if flag is not None:
+            after = [
+                ast.Assign(
+                    targets=[ast.Name(id=s.target.id, ctx=ast.Store())],
+                    value=ast.Name(id=kept, ctx=ast.Load()),
+                )
+            ]
+        if s.orelse and flag is not None:
+            # The for-else stores the DO variable's completion value, a
+            # Python int when the bounds are static; the carry it joins
+            # is the int32 index kept at the exit.
+            completion: list[ast.stmt] = []
+            for st in s.orelse:
+                if (
+                    isinstance(st, ast.Assign)
+                    and len(st.targets) == 1
+                    and isinstance(st.targets[0], ast.Name)
+                    and st.targets[0].id == s.target.id
+                ):
+                    st = ast.Assign(
+                        targets=st.targets,
+                        value=ast.Call(
+                            func=ast.Attribute(
+                                value=ast.Name(id="jnp", ctx=ast.Load()),
+                                attr="asarray",
+                                ctx=ast.Load(),
+                            ),
+                            args=[st.value],
+                            keywords=[
+                                ast.keyword(
+                                    arg="dtype",
+                                    value=ast.Attribute(
+                                        value=ast.Name(id="jnp", ctx=ast.Load()),
+                                        attr="int32",
+                                        ctx=ast.Load(),
+                                    ),
+                                )
+                            ],
+                        ),
+                    )
+                completion.append(st)
+            after += self.lower_block(
+                [ast.If(test=_not_flag(flag), body=completion, orelse=[])], depth
+            )
         if not body or all(isinstance(b, ast.Pass) for b in body):
             # A loop whose body lowered to nothing: statistics calls the
             # stand-in dropped, a nested loop that went the same way.
@@ -715,7 +926,7 @@ def lower_for(self, s, depth):
         if step == 1:
             fn = self._carry_fn(fname, [s.target.id, "_c"], carried, body)
             carry_call.args = [lo, hi, ast.Name(id=fname, ctx=ast.Load()), init]
-            return [fn, result]
+            return [*pre_flag, fn, result, *after]
 
         if stride is not None:
             lo_name, st_name, cnt_name = f"_lo_{self.n}", f"_st_{self.n}", f"_cnt_{self.n}"
@@ -754,7 +965,7 @@ def lower_for(self, s, depth):
                 ast.Name(id=fname, ctx=ast.Load()),
                 init,
             ]
-            return [*pre, fn, result]
+            return [*pre_flag, *pre, fn, result, *after]
 
         # step -1: Fortran DO k=hi,lo,-1 arrived as range(hi, stop, -1)
         # iterating hi..stop+1. Remap: t in [0, hi-stop), k = hi - t.
@@ -782,7 +993,7 @@ def lower_for(self, s, depth):
             ast.Name(id=fname, ctx=ast.Load()),
             init,
         ]
-        return [*pre, fn, result]
+        return [*pre_flag, *pre, fn, result, *after]
 
     def lower_if(self, s, depth):
         """if/elif/else -> lax.cond with the union of both branches'
@@ -907,7 +1118,14 @@ def split_params(fn_src):
 
 
 def emit_kernel(
-    fn_src, sub, closure, call_map=None, known_subs=None, writes=(), traced_scalars=frozenset()
+    fn_src,
+    sub,
+    closure,
+    call_map=None,
+    known_subs=None,
+    writes=(),
+    traced_scalars=frozenset(),
+    hosted=None,
 ):
     """Original numpy FunctionDef -> unparsed _<name>_k_impl source.
 
@@ -959,7 +1177,10 @@ def visit_FunctionDef(self, node):
             fn.body[at] = Extend().visit(stmt)
     _bind_writer_calls(fn, call_map or {})
     ExprMap().visit(fn)
-    CallRewrite(call_map or {}, known_subs or set()).visit(fn)
+    calls = CallRewrite(call_map or {}, known_subs or set())
+    calls.visit(fn)
+    if hosted is not None:
+        hosted.extend(calls.host_calls)
     nums, _ = static_spec(fn_src, sub, traced_scalars)
     required, _, _ = split_params(fn_src)
     statics = {required[pos] for pos in nums}
@@ -1058,16 +1279,17 @@ def static_spec(fn_src, sub, traced_scalars=frozenset()):
 
 def build_module(
     interface: dict[str, Any], tree: ast.Module, traced_scalars: frozenset[str] = frozenset()
-) -> tuple[list[str], list[str], dict[str, str]]:
+) -> tuple[list[str], list[str], dict[str, str], dict[str, list[str]]]:
     """Emit all kernels of one module to a fixpoint.
 
-    Returns (pieces, jitted, delegated) where pieces are source chunks
+    Returns (pieces, jitted, delegated, hosted) where pieces are source chunks
     (kernels + jit lines + wrappers + delegations + _JAX_KERNELS) and
     delegated maps name -> reason."""
     fns = {n.name: n for n in tree.body if isinstance(n, ast.FunctionDef)}
     subs = {s["name"]: s for s in interface["subprograms"]}
 
     delegated = {}
+    hosted: dict[str, list[str]] = {}
     kernels = set()
     for name, rec in subs.items():
         if name not in fns:
@@ -1103,6 +1325,7 @@ def build_module(
             }
         for name in sorted(emit_set):
             call_map = {k: v for k, v in call_map_all.items() if k != name}
+            calls_kept: list[str] = []
             try:
                 srcs[name] = emit_kernel(
                     fns[name],
@@ -1112,9 +1335,15 @@ def build_module(
                     set(subs),
                     wclosures[name],
                     traced_scalars,
+                    calls_kept,
                 )
             except JaxQueue as e:
                 failed[name] = f"[emit] {e}"
+                continue
+            if calls_kept:
+                hosted[name] = sorted(set(calls_kept))
+            else:
+                hosted.pop(name, None)
         if not failed:
             break
         for n, r in failed.items():
@@ -1188,7 +1417,7 @@ def build_module(
         if rec["name"] in delegated and rec["name"] in fns:
             pieces.append(f"{rec['name']} = _host.{rec['name']}")
     pieces.append(f"_JAX_KERNELS = {sorted(jitted)!r}")
-    return pieces, jitted, delegated
+    return pieces, jitted, delegated, hosted
 
 
 HEADER = '''"""Machine-generated by recast.transform.jax -- JAX backend (EXPERIMENTAL).
diff --git a/src/recast/transform/jax/runtime.py b/src/recast/transform/jax/runtime.py
index f14dfb6..6aad27b 100644
--- a/src/recast/transform/jax/runtime.py
+++ b/src/recast/transform/jax/runtime.py
@@ -174,8 +174,22 @@ def _f_fori(lo, hi, body, init):
     trace time. A dynamic bound is left to ``fori_loop``.
     """
     static = (int, np.integer)
-    if isinstance(lo, static) and isinstance(hi, static) and int(hi) <= int(lo):
-        return init
+    if isinstance(lo, static) and isinstance(hi, static):
+        if int(hi) <= int(lo):
+            return init
+        # A static trip count: a scan over the indices (reverse-
+        # differentiable, as fori_loop's own scan form is), spelled int32
+        # -- Fortran's default integer, the dtype every integer local and
+        # dummy carries, so a store of the index into one keeps its dtype
+        # across a lax.cond. fori_loop's scan form would count in the
+        # default int, int64 under x64.
+        indices = jnp.arange(int(lo), int(hi), dtype=jnp.int32)
+
+        def step(carry, i):
+            return body(i, carry), None
+
+        carry, _ = lax.scan(step, init, indices)
+        return carry
     return lax.fori_loop(lo, hi, body, init)
 
 
diff --git a/src/recast/transform/jax/translate.py b/src/recast/transform/jax/translate.py
index a70bbcb..f00d448 100644
--- a/src/recast/transform/jax/translate.py
+++ b/src/recast/transform/jax/translate.py
@@ -63,7 +63,7 @@ def apply(self, unit: Unit, facts: Facts, config: dict[str, Any]) -> Candidate:
 
         source = anchor.files[Path(f"{module}_numpy.py")].decode()
         tree = ast.parse(source)
-        pieces, jitted, delegated = build_module(facts.interface, tree)
+        pieces, jitted, delegated, hosted = build_module(facts.interface, tree)
         emitted = (
             HEADER.format(module=module, constants=constants_stem, runtime=runtime_stem)
             + _signatures_of(tree)
@@ -88,6 +88,9 @@ def apply(self, unit: Unit, facts: Facts, config: dict[str, Any]) -> Candidate:
                     "anchor": f"{module}_numpy.py",
                     "kernels": sorted(jitted),
                     "delegated": dict(sorted(delegated.items())),
+                    "host_calls": {
+                        k: [f"{module}.{n}" for n in v] for k, v in sorted(hosted.items())
+                    },
                     "runtime": f"{runtime_stem}.py",
                 },
             },
diff --git a/src/recast/transform/jax/tree.py b/src/recast/transform/jax/tree.py
index 6b50c04..e998f75 100644
--- a/src/recast/transform/jax/tree.py
+++ b/src/recast/transform/jax/tree.py
@@ -2950,7 +2950,12 @@ def apply(self, unit: Unit, facts: Facts, config: dict[str, Any]) -> Candidate:
         flat_tree, interface, flat_notes = flattened_module(
             tree, facts.interface, plans, ported, bundled
         )
-        pieces, jitted, delegated = build_module(interface, flat_tree, TRACED_SCALARS)
+        pieces, jitted, delegated, kept_on_host = build_module(interface, flat_tree, TRACED_SCALARS)
+        for name, kept in kept_on_host.items():
+            # This module's own subprograms the backend left on the host,
+            # beside the companions' the flat rewrite did.
+            own = flat_notes.setdefault("host_calls", {}).setdefault(name, [])
+            flat_notes["host_calls"][name] = sorted(set(own) | {f"{module}.{n}" for n in kept})
         # A flat function the backend delegated has a host to fall back on
         # only if the NumPy module carries its wrapper -- the gated ones. A
         # private subprogram's or a function's flat form has none, and a line
diff --git a/tests/test_jax_transform.py b/tests/test_jax_transform.py
index 07d1cfe..358db92 100644
--- a/tests/test_jax_transform.py
+++ b/tests/test_jax_transform.py
@@ -206,7 +206,7 @@ def test_a_module_state_write_threads_through_the_closure() -> None:
             },
         ]
     }
-    pieces, jitted, _delegated = build_module(interface, tree)
+    pieces, jitted, _delegated, _hosted = build_module(interface, tree)
     assert sorted(jitted) == ["tick", "use_tick"]
     text = "\n\n".join(pieces)
     assert "_host.cache = _res[0]" in text  # tick's wrapper stores the write back
@@ -276,24 +276,38 @@ def test_a_constant_table_is_read_through_jnp() -> None:
       end if
     end do
   end subroutine first_above
+  subroutine exit_index(n, zlo, z, kfound)
+    integer,  intent(in)  :: n
+    real(r8), intent(in)  :: zlo
+    real(r8), intent(in)  :: z(n)
+    integer,  intent(out) :: kfound
+    integer :: k
+    do k = 1, n
+      if ( z(k) > zlo ) exit
+    end do
+    kfound = k
+  end subroutine exit_index
 end module cycle_demo
 """
 
 
-def test_a_cycle_folds_into_the_branch_and_an_exit_is_delegated(tmp_path: Path) -> None:
+def test_a_cycle_folds_into_the_branch_and_an_exit_is_a_carried_flag(tmp_path: Path) -> None:
     """CLUBB's interpolators: ``if ( ... ) then ... cycle end if`` in a DO
     loop. The lowering passed the ``continue`` through into a ``lax.cond``
     branch -- a SyntaxError that took the whole emitted module down. Folded
-    into the branch structure it is a kernel; an EXIT has no fori_loop
-    shape and is delegated to the host, not emitted."""
+    into the branch structure it is a kernel. An EXIT is a flag the loop
+    carries: the trips after it do nothing, and the DO variable's value at
+    the exit is what the code after the loop reads (CLUBB's window search
+    and its sponge damping)."""
     import importlib
     import sys
 
     candidate = port(tmp_path, CYCLES, "cycle_demo")
-    assert candidate.notes["jax"]["kernels"] == ["clip_below"]
-    assert "first_above" in candidate.notes["jax"]["delegated"]
+    assert candidate.notes["jax"]["kernels"] == ["clip_below", "exit_index", "first_above"]
+    assert candidate.notes["jax"]["delegated"] == {}
     emitted = candidate.files[Path("cycle_demo_jax.py")].decode()
     assert "continue" not in emitted.replace("continuation", "")
+    assert "break" not in emitted
     out = tmp_path / "emitted"
     out.mkdir()
     for path, content in candidate.files.items():
@@ -308,6 +322,9 @@ def test_a_cycle_folds_into_the_branch_and_an_exit_is_delegated(tmp_path: Path)
         w = np.asarray(module.clip_below(4, 2.5, z, v))
         assert w.tolist() == [0.0, 0.0, 2.0, 2.0]
         assert int(module.first_above(4, 2.5, z)) == 3
+        assert int(module.first_above(4, 9.0, z)) == 0
+        assert int(module.exit_index(4, 2.5, z)) == 3
+        assert int(module.exit_index(4, 9.0, z)) == 5  # ran to completion: n + 1
     finally:
         sys.path.remove(str(out))
         for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"):
@@ -1549,6 +1566,67 @@ def test_a_callee_extent_argument_is_the_callers_axis(tmp_path: Path) -> None:
             sys.modules.pop(f"sized_mod{suffix}", None)
 
 
+OWN_CHECK_UNDER_A_SWITCH = """\
+module ownswitch_mod
+  implicit none
+contains
+  subroutine complain( n, x, msg )
+    integer, intent(in) :: n
+    real(8), intent(in) :: x(n)
+    character(len=*), intent(out) :: msg
+    msg = "fine"
+    if ( any( x < 0.0d0 ) ) msg = "negative"
+  end subroutine complain
+  subroutine step( n, l_check, x, y )
+    integer, intent(in) :: n
+    logical, intent(in) :: l_check
+    real(8), intent(in) :: x(n)
+    real(8), intent(out) :: y(n)
+    character(len=16) :: msg
+    y = 2.0d0 * x
+    if ( l_check ) then
+      call complain( n, y, msg )
+    end if
+  end subroutine step
+end module ownswitch_mod
+"""
+
+
+def test_a_subprogram_of_the_module_the_port_could_not_emit_stays_on_the_host(
+    tmp_path: Path,
+) -> None:
+    """pdf_closure_driver calls its zm variant under a switch the run has
+    off; fill_holes_vertical_api dispatches on a type. A same-module
+    callee the port could not emit no longer delegates its caller and the
+    callers above it: the call stays the host's (``_host.<name>``) under
+    the anchor's guard, named in the notes, never traced while the guard
+    holds."""
+    import importlib
+    import sys
+
+    candidate = port(tmp_path, OWN_CHECK_UNDER_A_SWITCH, "ownswitch_mod")
+    assert candidate.notes["jax"]["kernels"] == ["step"], candidate.notes["jax"]
+    assert "complain" in candidate.notes["jax"]["delegated"]
+    assert candidate.notes["jax"]["host_calls"] == {"step": ["ownswitch_mod.complain"]}
+    emitted = candidate.files[Path("ownswitch_mod_jax.py")].decode()
+    assert "_host.complain(" in emitted
+    out = tmp_path / "emitted"
+    out.mkdir()
+    for path, content in candidate.files.items():
+        (out / path.name).write_bytes(content)
+    sys.path.insert(0, str(out))
+    try:
+        module = importlib.import_module("ownswitch_mod_jax")
+        import numpy as np
+
+        got = module.step(2, False, np.array([1.0, -1.0]))
+        assert np.asarray(got).tolist() == [2.0, -2.0]
+    finally:
+        sys.path.remove(str(out))
+        for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"):
+            sys.modules.pop(f"ownswitch_mod{suffix}", None)
+
+
 SILENT_CHECK = """\
 module silent_mod
   implicit none

From 60d596c8da00151f813e3fd98c4b1102ca3e2a3f Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Sat, 5 Sep 2026 02:04:40 -0600
Subject: [PATCH 67/73] jax: a window with traced bounds a static distance
 apart is a gather

A slice whose bounds trace (the loop index, through locals assigned
k -+ dir * n in scope) but whose length does not becomes the index
lo + arange(trips) * step -- a gather of static length at a traced
offset, a scatter when stored -- with the trip count the runtime's
_f_trips over static atoms, a Python int at trace time. The mask rules
keep the slices whose length depends on a traced bound.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/jax/tree.py | 220 ++++++++++++++++++++++++++++++-
 tests/test_jax_transform.py      |  62 +++++++++
 2 files changed, 281 insertions(+), 1 deletion(-)

diff --git a/src/recast/transform/jax/tree.py b/src/recast/transform/jax/tree.py
index e998f75..6938774 100644
--- a/src/recast/transform/jax/tree.py
+++ b/src/recast/transform/jax/tree.py
@@ -1667,9 +1667,13 @@ def _rewritten_body(fn: ast.FunctionDef, rewrite: _Rewrite) -> list[ast.stmt]:
     rewrite.inits = _guard_inits(fn)
     rewrite.statics = frozenset(rewrite.statics) | _concrete_scalars(fn, rewrite)
     lowered: list[ast.stmt] = []
+    body = [copy.deepcopy(s) for s in fn.body]
+    # Fixed-length windows with traced bounds first, so the mask rules never
+    # see them: ``x(k-w:k+w)`` is a gather at ``lo + arange(2w+1)``.
+    _window_slices(body, frozenset(rewrite.statics))
     # Single exit first, on the anchor's own returns: the flat return that
     # replaces them is one statement at the end.
-    for statement in _single_exit([copy.deepcopy(s) for s in fn.body]):
+    for statement in _single_exit(body):
         result = rewrite.visit(statement)
         if result is None:
             continue
@@ -1759,6 +1763,220 @@ def _rebuilt_objects(lowered: list[ast.stmt], rewrite: _Rewrite) -> list[ast.stm
     return rebuilt
 
 
+Affine = tuple[dict[str, int], int]
+"""A linear form: ``{name: coefficient}`` and a constant."""
+
+
+def _static_name(name: str, statics: frozenset[str]) -> bool:
+    return name in statics or (name.isupper() and len(name) > 1)
+
+
+def _static_atom(node: ast.expr, statics: frozenset[str]) -> bool:
+    """An expression over static names and constants only -- ``dir * n``,
+    a Python int at trace time -- taken whole as one term of a form."""
+    names = [n for n in ast.walk(node) if isinstance(n, ast.Name)]
+    return bool(names) and all(_static_name(n.id, statics) for n in names)
+
+
+def _affine(
+    node: ast.expr, env: dict[str, ast.expr | None], statics: frozenset[str], depth: int = 0
+) -> Affine | None:
+    """The linear form of an integer expression: ``{term: coefficient}``
+    and a constant, a term being a traced name or a static atom (spelled
+    ``@<text>``, over static names only). A local that was assigned an
+    affine expression in scope (``k_start = k - dir * n``) is
+    substituted; None where the expression is not linear."""
+    if depth > 8:
+        return None
+    if isinstance(node, ast.Constant):
+        if isinstance(node.value, int) and not isinstance(node.value, bool):
+            return {}, int(node.value)
+        return None
+    if isinstance(node, ast.Name):
+        bound = env.get(node.id)
+        if bound is not None:
+            inner = _affine(bound, env, statics, depth + 1)
+            if inner is not None:
+                return inner
+        if _static_name(node.id, statics):
+            return {f"@{node.id}": 1}, 0
+        return {node.id: 1}, 0
+    if not isinstance(node, ast.Name | ast.Constant) and _static_atom(node, statics):
+        return {f"@{ast.unparse(node)}": 1}, 0
+    if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.UAdd | ast.USub):
+        inner = _affine(node.operand, env, statics, depth)
+        if inner is None:
+            return None
+        if isinstance(node.op, ast.UAdd):
+            return inner
+        return {k: -v for k, v in inner[0].items()}, -inner[1]
+    if isinstance(node, ast.BinOp) and isinstance(node.op, ast.Add | ast.Sub | ast.Mult):
+        left = _affine(node.left, env, statics, depth)
+        right = _affine(node.right, env, statics, depth)
+        if left is None or right is None:
+            return None
+        if isinstance(node.op, ast.Mult):
+            if not left[0]:
+                scale, form = left[1], right
+            elif not right[0]:
+                scale, form = right[1], left
+            else:
+                return None
+            return {k: v * scale for k, v in form[0].items()}, form[1] * scale
+        sign = 1 if isinstance(node.op, ast.Add) else -1
+        coeffs = dict(left[0])
+        for k, v in right[0].items():
+            coeffs[k] = coeffs.get(k, 0) + sign * v
+        return {k: v for k, v in coeffs.items() if v}, left[1] + sign * right[1]
+    if (
+        isinstance(node, ast.Call)
+        and len(node.args) == 1
+        and not node.keywords
+        and (
+            (isinstance(node.func, ast.Name) and node.func.id == "int")
+            or (
+                isinstance(node.func, ast.Attribute)
+                and node.func.attr in ("int32", "int64")
+                and isinstance(node.func.value, ast.Name)
+                and node.func.value.id in ("np", "jnp")
+            )
+        )
+    ):
+        return _affine(node.args[0], env, statics, depth)
+    return None
+
+
+def _affine_ast(form: Affine) -> ast.expr:
+    """``2 * (dir * n) + 1`` back as an expression, over static atoms only."""
+    coeffs, const = form
+    out: ast.expr = ast.Constant(value=const)
+    for key in sorted(coeffs):
+        atom: ast.expr = ast.parse(key[1:], mode="eval").body
+        term: ast.expr = ast.BinOp(left=ast.Constant(value=coeffs[key]), op=ast.Mult(), right=atom)
+        out = ast.BinOp(left=out, op=ast.Add(), right=term)
+    return out
+
+
+class _Windows(ast.NodeTransformer):
+    """``x[i - 1, k_start - 1:k_end:dir]`` where ``k_start`` and ``k_end``
+    are ``k -+ dir * n``: bounds that trace, a distance apart that does not.
+    The slice becomes the index ``lo + arange(trips) * step`` -- a gather
+    of static length at a traced offset, a scatter when stored -- with
+    the trip count the runtime's ``_f_trips(0, distance, step)`` over
+    static names, a Python int at trace time. Fortran's windows (CLUBB's
+    hole filler draws ``num_hf_draw_points`` levels either side of ``k``)
+    are this; a slice whose length depends on a traced bound is the mask
+    rules' business and is left to them."""
+
+    def __init__(self, env: dict[str, ast.expr | None], statics: frozenset[str]) -> None:
+        self.env = env
+        self.statics = statics
+
+    def _dynamic(self, node: ast.expr) -> bool:
+        return any(
+            isinstance(n, ast.Name) and not _static_name(n.id, self.statics) for n in ast.walk(node)
+        )
+
+    def _window(self, element: ast.Slice) -> ast.expr | None:
+        if element.upper is None:
+            return None
+        lower: ast.expr = element.lower if element.lower is not None else ast.Constant(0)
+        if not (self._dynamic(lower) or self._dynamic(element.upper)):
+            return None
+        if element.step is not None and self._dynamic(element.step):
+            return None
+        hi = _affine(element.upper, self.env, self.statics)
+        lo = _affine(lower, self.env, self.statics)
+        if hi is None or lo is None:
+            return None
+        coeffs = dict(hi[0])
+        for k, v in lo[0].items():
+            coeffs[k] = coeffs.get(k, 0) - v
+        coeffs = {k: v for k, v in coeffs.items() if v}
+        if any(not k.startswith("@") for k in coeffs):
+            return None  # a traced term survives: the length traces too
+        distance = _affine_ast((coeffs, hi[1] - lo[1]))
+        step: ast.expr = (
+            copy.deepcopy(element.step) if element.step is not None else ast.Constant(1)
+        )
+        trips = ast.Call(
+            func=ast.Name(id="_f_trips", ctx=ast.Load()),
+            args=[ast.Constant(0), distance, step],
+            keywords=[],
+        )
+        offsets: ast.expr = _jnp("arange", [trips])
+        if not (isinstance(step, ast.Constant) and step.value == 1):
+            offsets = ast.BinOp(left=offsets, op=ast.Mult(), right=copy.deepcopy(step))
+        return ast.BinOp(left=copy.deepcopy(lower), op=ast.Add(), right=offsets)
+
+    def visit_Subscript(self, node: ast.Subscript) -> ast.AST:
+        self.generic_visit(node)
+        elts = list(node.slice.elts) if isinstance(node.slice, ast.Tuple) else [node.slice]
+        changed = False
+        out: list[ast.expr] = []
+        for element in elts:
+            window = self._window(element) if isinstance(element, ast.Slice) else None
+            if window is None:
+                out.append(element)
+            else:
+                out.append(window)
+                changed = True
+        if changed:
+            node.slice = ast.Tuple(elts=out, ctx=ast.Load()) if len(out) > 1 else out[0]
+        return node
+
+
+def _window_slices(stmts: list[ast.stmt], statics: frozenset[str]) -> None:
+    """Rewrite the windows in a block in place, an affine local's
+    definition in scope where a bound names one."""
+
+    def block(body: list[ast.stmt], env: dict[str, ast.expr | None]) -> set[str]:
+        assigned: set[str] = set()
+        for st in body:
+            if isinstance(st, ast.Assign):
+                _Windows(env, statics).visit(st)
+                for target in st.targets:
+                    if isinstance(target, ast.Name):
+                        affine = _affine(st.value, env, statics) is not None
+                        env[target.id] = st.value if affine else None
+                        assigned.add(target.id)
+            elif isinstance(st, ast.AugAssign):
+                _Windows(env, statics).visit(st)
+                if isinstance(st.target, ast.Name):
+                    env[st.target.id] = None
+                    assigned.add(st.target.id)
+            elif isinstance(st, ast.If | ast.While):
+                _Windows(env, statics).visit(st.test)
+                inner = block(st.body, dict(env)) | block(st.orelse, dict(env))
+                for name in inner:
+                    env[name] = None
+                assigned |= inner
+            elif isinstance(st, ast.For):
+                _Windows(env, statics).visit(st.iter)
+                scope = dict(env)
+                if isinstance(st.target, ast.Name):
+                    scope[st.target.id] = None
+                inner = block(st.body, scope) | block(st.orelse, dict(scope))
+                if isinstance(st.target, ast.Name):
+                    inner.add(st.target.id)
+                for name in inner:
+                    env[name] = None
+                assigned |= inner
+            elif isinstance(st, ast.Try):
+                inner = block(st.body, dict(env))
+                for handler in st.handlers:
+                    inner |= block(handler.body, dict(env))
+                inner |= block(st.orelse, dict(env)) | block(st.finalbody, dict(env))
+                for name in inner:
+                    env[name] = None
+                assigned |= inner
+            else:
+                _Windows(env, statics).visit(st)
+        return assigned
+
+    block(stmts, {})
+
+
 def _dim_sources(args: list[dict[str, Any]]) -> dict[str, tuple[str, int]]:
     """``{extent name: (array dummy, axis)}`` for every dummy array whose
     declared extent is a bare name -- ``a(n)`` says ``n == a.shape[0]``."""
diff --git a/tests/test_jax_transform.py b/tests/test_jax_transform.py
index 358db92..e02eac9 100644
--- a/tests/test_jax_transform.py
+++ b/tests/test_jax_transform.py
@@ -1627,6 +1627,68 @@ def test_a_subprogram_of_the_module_the_port_could_not_emit_stays_on_the_host(
             sys.modules.pop(f"ownswitch_mod{suffix}", None)
 
 
+WINDOWED = """\
+module window_mod
+  implicit none
+contains
+  subroutine fill( n, w, dir, x, y )
+    integer, intent(in) :: n, w, dir
+    real(8), intent(in) :: x(n)
+    real(8), intent(out) :: y(n)
+    integer :: k, lo, hi
+    y = x
+    do k = 1 + w, n - w
+      lo = k - dir * w
+      hi = k + dir * w
+      if ( any( x(lo:hi:dir) < 0.0d0 ) ) then
+        y(k) = -1.0d0
+      else
+        y(k) = maxval( x(lo:hi:dir) )
+      end if
+    end do
+  end subroutine fill
+end module window_mod
+"""
+
+
+def test_a_window_with_traced_bounds_a_static_distance_apart_is_a_gather(tmp_path: Path) -> None:
+    """CLUBB's sliding-window hole filler reads ``field(i, k_start:k_end:dir)``
+    with ``k_start = k - dir * n`` and ``k_end = k + dir * n``: bounds that
+    trace (the loop index), a distance apart that does not. The slice is a
+    gather at ``lo + arange(trips) * step`` -- the mask rules, for a slice
+    whose length depends on a traced bound, never see it."""
+    import importlib
+    import sys
+
+    from recast.transform.jax.tree import TreeToJax
+    from recast.transform.numpy.tree import TreeConventions
+
+    (tmp_path / "window_mod.f90").write_text(WINDOWED)
+    frontend = FortranFrontend(flatten=True)
+    unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:window_mod")
+    facts = frontend.analyze(unit, tmp_path)
+    candidate = TreeToJax(TreeConventions()).apply(unit, facts, {"root": str(tmp_path)})
+    assert candidate.notes["jax"]["kernels"] == ["fill"], candidate.notes["jax"]
+    emitted = candidate.files[Path("window_mod_jax.py")].decode()
+    assert "jnp.arange(_f_trips(0, " in emitted, emitted
+    out = tmp_path / "emitted"
+    out.mkdir()
+    for path, content in candidate.files.items():
+        (out / path.name).write_bytes(content)
+    sys.path.insert(0, str(out))
+    try:
+        module = importlib.import_module("window_mod_jax")
+        import numpy as np
+
+        x = np.array([1.0, 5.0, 2.0, -3.0, 4.0, 6.0])
+        got = np.asarray(module.fill(6, 1, 1, x)).tolist()
+        assert got == [1.0, 5.0, -1.0, -1.0, -1.0, 6.0]
+    finally:
+        sys.path.remove(str(out))
+        for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"):
+            sys.modules.pop(f"window_mod{suffix}", None)
+
+
 SILENT_CHECK = """\
 module silent_mod
   implicit none

From 162a507b64cc5849e822943a4ad6f7f06f63ee8b Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Sat, 5 Sep 2026 02:08:48 -0600
Subject: [PATCH 68/73] jax: an array passed through a reshape comes back in
 its own shape

Sequence association: the anchor spells a 2-d right-hand side handed to
a solver declared over three axes as np.reshape(rhs, (n, m, 1),
order='F'). The kernel's output for that dummy is reshaped to the
array's own shape, the same order, and rebinds it -- a reshape is no
store target, and the anchor's copy-out spelling does not change that.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/jax/tree.py | 65 ++++++++++++++++++++++++++++++--
 tests/test_jax_transform.py      | 60 +++++++++++++++++++++++++++++
 2 files changed, 122 insertions(+), 3 deletions(-)

diff --git a/src/recast/transform/jax/tree.py b/src/recast/transform/jax/tree.py
index 6938774..1b24fa2 100644
--- a/src/recast/transform/jax/tree.py
+++ b/src/recast/transform/jax/tree.py
@@ -1413,6 +1413,12 @@ def _rewrite_call(self, target: ast.expr | None, call: ast.Call) -> Any:
                 targets.append(ast.Name(id=name, ctx=ast.Store()))
             else:
                 where = slot.get(name)
+                passed_as = actual_by_dummy.get(name.lower())
+                if passed_as is not None and _reshaped_name(passed_as) is not None:
+                    # The actual is a reshaped view of a name (sequence
+                    # association); the output comes back in the callee's
+                    # shape, whatever the anchor's own target spelling.
+                    where = passed_as
                 if where is None:
                     raise NotFlat(f"{callee.subprogram['name']}: output {name} has no target")
                 bound: Any = self.visit(copy.deepcopy(where))
@@ -1422,9 +1428,38 @@ def _rewrite_call(self, target: ast.expr | None, call: ast.Call) -> Any:
                     self.temps += 1
                     temp = f"_t{self.temps}"
                     targets.append(ast.Name(id=temp, ctx=ast.Store()))
-                    follow.append(
-                        ast.Assign(targets=[bound], value=ast.Name(id=temp, ctx=ast.Load()))
-                    )
+                    reshaped = _reshaped_name(bound)
+                    if reshaped is not None:
+                        # ``np.reshape(rhs, (n, m, 1), order='F')`` as the
+                        # actual: sequence association, the anchor's view of
+                        # the array in the callee's shape. What comes back
+                        # is reshaped to the array's own shape, the same
+                        # order, and rebinds it.
+                        follow.append(
+                            ast.Assign(
+                                targets=[ast.Name(id=reshaped, ctx=ast.Store())],
+                                value=ast.Call(
+                                    func=ast.Attribute(
+                                        value=ast.Name(id="jnp", ctx=ast.Load()),
+                                        attr="reshape",
+                                        ctx=ast.Load(),
+                                    ),
+                                    args=[
+                                        ast.Name(id=temp, ctx=ast.Load()),
+                                        ast.Attribute(
+                                            value=ast.Name(id=reshaped, ctx=ast.Load()),
+                                            attr="shape",
+                                            ctx=ast.Load(),
+                                        ),
+                                    ],
+                                    keywords=[ast.keyword(arg="order", value=ast.Constant("F"))],
+                                ),
+                            )
+                        )
+                    else:
+                        follow.append(
+                            ast.Assign(targets=[bound], value=ast.Name(id=temp, ctx=ast.Load()))
+                        )
         if not targets:
             return ast.Expr(value=new_call)
         # One output comes back bare (the flat function returns a name, not
@@ -1767,6 +1802,30 @@ def _rebuilt_objects(lowered: list[ast.stmt], rewrite: _Rewrite) -> list[ast.stm
 """A linear form: ``{name: coefficient}`` and a constant."""
 
 
+def _reshaped_name(node: ast.expr) -> str | None:
+    """``np.reshape(x, shape, order='F')`` / ``jnp.reshape(...)`` /
+    ``x.reshape(...)`` over a bare name: that name."""
+    if not isinstance(node, ast.Call):
+        return None
+    func = node.func
+    if (
+        isinstance(func, ast.Attribute)
+        and func.attr == "reshape"
+        and isinstance(func.value, ast.Name)
+        and func.value.id in ("np", "jnp")
+        and node.args
+        and isinstance(node.args[0], ast.Name)
+    ):
+        return node.args[0].id
+    if (
+        isinstance(func, ast.Attribute)
+        and func.attr == "reshape"
+        and isinstance(func.value, ast.Name)
+    ):
+        return func.value.id
+    return None
+
+
 def _static_name(name: str, statics: frozenset[str]) -> bool:
     return name in statics or (name.isupper() and len(name) > 1)
 
diff --git a/tests/test_jax_transform.py b/tests/test_jax_transform.py
index e02eac9..fb003b3 100644
--- a/tests/test_jax_transform.py
+++ b/tests/test_jax_transform.py
@@ -1689,6 +1689,66 @@ def test_a_window_with_traced_bounds_a_static_distance_apart_is_a_gather(tmp_pat
             sys.modules.pop(f"window_mod{suffix}", None)
 
 
+RESHAPED_ACTUAL = """\
+module reshaped_mod
+  implicit none
+  type knobs_type
+    real(8) :: gain = 2.0d0
+  end type knobs_type
+contains
+  subroutine solve( n, k, rhs )
+    integer, intent(in) :: n
+    type(knobs_type), intent(in) :: k
+    real(8), intent(inout) :: rhs(n, 1)
+    rhs(:, 1) = k%gain * rhs(:, 1)
+  end subroutine solve
+  subroutine step( n, k, x )
+    integer, intent(in) :: n
+    type(knobs_type), intent(in) :: k
+    real(8), intent(inout) :: x(n)
+    call solve( n, k, x )
+    x = x + 1.0d0
+  end subroutine step
+end module reshaped_mod
+"""
+
+
+def test_an_array_passed_through_a_reshape_comes_back_in_its_own_shape(tmp_path: Path) -> None:
+    """CLUBB hands a 2-d right-hand side to a solver declared over three
+    axes: sequence association, which the anchor spells as
+    ``np.reshape(rhs, (n, m, 1), order='F')`` for the actual. The kernel's
+    output is reshaped back to the array's own shape and rebinds it -- a
+    reshape is no store target."""
+    import importlib
+    import sys
+
+    from recast.transform.jax.tree import TreeToJax
+    from recast.transform.numpy.tree import TreeConventions
+
+    (tmp_path / "reshaped_mod.f90").write_text(RESHAPED_ACTUAL)
+    frontend = FortranFrontend(flatten=True)
+    unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:reshaped_mod")
+    facts = frontend.analyze(unit, tmp_path)
+    candidate = TreeToJax(TreeConventions()).apply(unit, facts, {"root": str(tmp_path)})
+    assert "reshape(x" in candidate.files[Path("reshaped_mod_numpy.py")].decode()
+    assert "step_flat" in candidate.notes["jax"]["kernels"], candidate.notes["jax"]
+    out = tmp_path / "emitted"
+    out.mkdir()
+    for path, content in candidate.files.items():
+        (out / path.name).write_bytes(content)
+    sys.path.insert(0, str(out))
+    try:
+        module = importlib.import_module("reshaped_mod_jax")
+        import numpy as np
+
+        got = module.step_flat(2, np.array([1.0, 3.0]), 2, np.float64(2.0))
+        assert np.asarray(got).tolist() == [3.0, 7.0]
+    finally:
+        sys.path.remove(str(out))
+        for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"):
+            sys.modules.pop(f"reshaped_mod{suffix}", None)
+
+
 SILENT_CHECK = """\
 module silent_mod
   implicit none

From 7d6c987480130b1c307b6a371cbb1446a11051ea Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Sat, 5 Sep 2026 02:29:27 -0600
Subject: [PATCH 69/73] jax: a static branch whose arms lowered to nothing is
 no branch

A debug print under a level check, statistics under their switch: the
arms lower to nothing and the static Python if was emitted with no
body, a SyntaxError that took the whole emitted module down. Nothing in
both arms is no branch; nothing in one arm is pass.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/jax/backend.py |  6 ++++
 tests/test_jax_transform.py         | 49 +++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+)

diff --git a/src/recast/transform/jax/backend.py b/src/recast/transform/jax/backend.py
index 14847d3..972ea9d 100644
--- a/src/recast/transform/jax/backend.py
+++ b/src/recast/transform/jax/backend.py
@@ -1008,6 +1008,12 @@ def lower_if(self, s, depth):
         a None arg — is never traced)."""
         body = self.lower_block(s.body, depth + 1)
         orelse = self.lower_block(s.orelse, depth + 1)
+        if not body and not orelse:
+            # Both arms lowered to nothing (a debug print under a level
+            # check, statistics under their switch): no branch at all.
+            return []
+        if not body:
+            body = [ast.Pass()]  # a Python if needs a body; the else is the point
         if _static_test(s.test):
             return [ast.If(test=s.test, body=body, orelse=orelse or [])]
         if _static_test(s.test, self.statics):
diff --git a/tests/test_jax_transform.py b/tests/test_jax_transform.py
index fb003b3..bbb6496 100644
--- a/tests/test_jax_transform.py
+++ b/tests/test_jax_transform.py
@@ -1749,6 +1749,55 @@ def test_an_array_passed_through_a_reshape_comes_back_in_its_own_shape(tmp_path:
             sys.modules.pop(f"reshaped_mod{suffix}", None)
 
 
+EMPTY_GUARD = """\
+module emptyguard_mod
+  implicit none
+contains
+  subroutine scale( n, level, x, y )
+    integer, intent(in) :: n, level
+    real(8), intent(in) :: x(n)
+    real(8), intent(out) :: y(n)
+    y = 2.0d0 * x
+    if ( level > 0 ) then
+      print *, "scaled", n
+    end if
+    if ( level > 1 ) then
+      print *, "twice"
+    else
+      y = y + 1.0d0
+    end if
+  end subroutine scale
+end module emptyguard_mod
+"""
+
+
+def test_a_static_branch_whose_arms_lowered_to_nothing_is_no_branch(tmp_path: Path) -> None:
+    """``if ( clubb_at_least_debug_level( 0 ) ) then`` around a print: the
+    print is dropped, the static Python if was emitted with no body -- a
+    SyntaxError that took the whole emitted module down. Nothing in both
+    arms is no branch; nothing in one arm is ``pass``."""
+    import importlib
+    import sys
+
+    candidate = port(tmp_path, EMPTY_GUARD, "emptyguard_mod")
+    assert candidate.notes["jax"]["kernels"] == ["scale"], candidate.notes["jax"]
+    out = tmp_path / "emitted"
+    out.mkdir()
+    for path, content in candidate.files.items():
+        (out / path.name).write_bytes(content)
+    sys.path.insert(0, str(out))
+    try:
+        module = importlib.import_module("emptyguard_mod_jax")
+        import numpy as np
+
+        assert np.asarray(module.scale(2, 0, np.array([1.0, 3.0]))).tolist() == [3.0, 7.0]
+        assert np.asarray(module.scale(2, 2, np.array([1.0, 3.0]))).tolist() == [2.0, 6.0]
+    finally:
+        sys.path.remove(str(out))
+        for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"):
+            sys.modules.pop(f"emptyguard_mod{suffix}", None)
+
+
 SILENT_CHECK = """\
 module silent_mod
   implicit none

From 423187f809693bac5d3976285f28599b3f38f36a Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Sat, 5 Sep 2026 02:57:49 -0600
Subject: [PATCH 70/73] jax: a module constant through its alias is static

iipdf_type == _mod.IIPDF_ADG1: pdf_closure picks its PDF by a constant
of model_flags spelled through the module alias. The static-test
predicate (and the tree's static-expression one) took only the bare
upper-case spelling; the dispatch was a lax.cond chain, every arm
traced, and the arm the run never takes carried a host driver's record
object into the cond. Static now, the dead arm is never traced.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/jax/backend.py | 15 ++++++
 src/recast/transform/jax/tree.py    |  8 +++
 tests/test_jax_transform.py         | 76 +++++++++++++++++++++++++++++
 3 files changed, 99 insertions(+)

diff --git a/src/recast/transform/jax/backend.py b/src/recast/transform/jax/backend.py
index 972ea9d..10625c9 100644
--- a/src/recast/transform/jax/backend.py
+++ b/src/recast/transform/jax/backend.py
@@ -495,6 +495,19 @@ def add(n):
     return out
 
 
+def _module_constant(node) -> bool:
+    """``_mod.IIPDF_ADG1``: a use-associated module constant through its
+    module alias -- a Python value at trace time, like the bare upper-case
+    spelling of one the translation resolved."""
+    return (
+        isinstance(node, ast.Attribute)
+        and isinstance(node.value, ast.Name)
+        and node.value.id.startswith("_")
+        and node.attr.isupper()
+        and len(node.attr) > 1
+    )
+
+
 def _static_test(test, statics=frozenset()):
     """True for branch conditions decidable at trace time per the
     translate.py grammar: `x is [not] None` (Fortran PRESENT) and bare
@@ -518,6 +531,8 @@ def constant(node):
             return isinstance(node.value, (int, float)) and not isinstance(node.value, bool)
         if isinstance(node, ast.Name):
             return node.id in statics or (node.id.isupper() and len(node.id) > 1)
+        if _module_constant(node):
+            return True
         if isinstance(node, ast.BinOp):
             return constant(node.left) and constant(node.right)
         if isinstance(node, ast.UnaryOp):
diff --git a/src/recast/transform/jax/tree.py b/src/recast/transform/jax/tree.py
index 1b24fa2..076037e 100644
--- a/src/recast/transform/jax/tree.py
+++ b/src/recast/transform/jax/tree.py
@@ -2787,6 +2787,14 @@ def _static_expression(node: ast.expr, statics: frozenset[str]) -> bool:
         return isinstance(node.value, (int, float)) and not isinstance(node.value, bool)
     if isinstance(node, ast.Name):
         return node.id.isupper() or node.id in statics
+    if (
+        isinstance(node, ast.Attribute)
+        and isinstance(node.value, ast.Name)
+        and node.value.id.startswith("_")
+        and node.attr.isupper()
+        and len(node.attr) > 1
+    ):
+        return True  # ``_mod.IIPDF_ADG1``: a module constant through its alias
     if isinstance(node, ast.BinOp) and isinstance(
         node.op, (ast.Add, ast.Sub, ast.Mult, ast.Div, ast.FloorDiv)
     ):
diff --git a/tests/test_jax_transform.py b/tests/test_jax_transform.py
index bbb6496..a21ab05 100644
--- a/tests/test_jax_transform.py
+++ b/tests/test_jax_transform.py
@@ -1798,6 +1798,82 @@ def test_a_static_branch_whose_arms_lowered_to_nothing_is_no_branch(tmp_path: Pa
             sys.modules.pop(f"emptyguard_mod{suffix}", None)
 
 
+PDF_KINDS = """\
+module kinds_mod
+  implicit none
+  integer, parameter :: I_PLAIN = 1
+  integer, parameter :: I_TWICE = 2
+  logical :: l_noisy = .false.
+end module kinds_mod
+"""
+
+DISPATCHES_ON_A_KIND = """\
+module dispatch_mod
+  use kinds_mod, only: I_PLAIN, I_TWICE
+  use silent_mod, only: complain
+  implicit none
+contains
+  subroutine pick( n, kind, x, y )
+    integer, intent(in) :: n, kind
+    real(8), intent(in) :: x(n)
+    real(8), intent(out) :: y(n)
+    if ( kind == I_TWICE ) then
+      y = 2.0d0 * x
+    else if ( kind == I_PLAIN ) then
+      y = x
+    else
+      y = 0.0d0
+      call complain( n, y )
+    end if
+  end subroutine pick
+end module dispatch_mod
+"""
+
+
+def test_a_dispatch_on_a_module_constant_through_its_alias_is_static(tmp_path: Path) -> None:
+    """``iipdf_type == _mod.IIPDF_ADG1``: pdf_closure picks its PDF by a
+    constant of model_flags, spelled through the module alias. Static, the
+    way the bare upper-case spelling is, so the arm the run never takes --
+    with its host call -- is never traced."""
+    import importlib
+    import sys
+
+    from recast.transform.jax.tree import TreeToJax
+    from recast.transform.numpy.tree import TreeConventions
+
+    (tmp_path / "kinds_mod.f90").write_text(PDF_KINDS)
+    (tmp_path / "silent_mod.f90").write_text(SILENT_CHECK)
+    (tmp_path / "dispatch_mod.f90").write_text(DISPATCHES_ON_A_KIND)
+    frontend = FortranFrontend(flatten=True)
+    unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:dispatch_mod")
+    facts = frontend.analyze(unit, tmp_path)
+    candidate = TreeToJax(TreeConventions()).apply(unit, facts, {"root": str(tmp_path)})
+    assert "pick" in candidate.notes["jax"]["kernels"], candidate.notes["jax"]
+    emitted = candidate.files[Path("dispatch_mod_jax.py")].decode()
+    anchor_src = candidate.files[Path("dispatch_mod_numpy.py")].decode()
+    assert "lax.cond(kind ==" not in emitted, emitted
+    out = tmp_path / "emitted"
+    out.mkdir()
+    for path, content in candidate.files.items():
+        (out / path.name).write_bytes(content)
+    sys.path.insert(0, str(out))
+    try:
+        module = importlib.import_module("dispatch_mod_jax")
+        import numpy as np
+
+        x = np.array([1.0, 3.0])
+        assert np.asarray(module.pick(2, 2, x)).tolist() == [2.0, 6.0]
+        assert np.asarray(module.pick(2, 1, x)).tolist() == [1.0, 3.0]
+        print(
+            "ANCHOR_SPELLING", [line for line in anchor_src.splitlines() if "I_TWICE" in line][:2]
+        )
+    finally:
+        sys.path.remove(str(out))
+        for name in list(sys.modules):
+            if name.startswith(("dispatch_mod", "kinds_mod", "silent_mod")):
+                sys.modules.pop(name, None)
+
+
 SILENT_CHECK = """\
 module silent_mod
   implicit none

From 316f36fe9f2d4bb2126614d40a61c98cccc8adcf Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Sat, 5 Sep 2026 02:59:00 -0600
Subject: [PATCH 71/73] tests: the dispatch test asserts the dual form, not the
 absence of its fallback

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 tests/test_jax_transform.py | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/tests/test_jax_transform.py b/tests/test_jax_transform.py
index a21ab05..28689bb 100644
--- a/tests/test_jax_transform.py
+++ b/tests/test_jax_transform.py
@@ -1850,8 +1850,9 @@ def test_a_dispatch_on_a_module_constant_through_its_alias_is_static(tmp_path: P
     candidate = TreeToJax(TreeConventions()).apply(unit, facts, {"root": str(tmp_path)})
     assert "pick" in candidate.notes["jax"]["kernels"], candidate.notes["jax"]
     emitted = candidate.files[Path("dispatch_mod_jax.py")].decode()
-    anchor_src = candidate.files[Path("dispatch_mod_numpy.py")].decode()
-    assert "lax.cond(kind ==" not in emitted, emitted
+    # The dual form: a Python if when the kind is concrete (through the jit
+    # wrapper it is), the lax.cond only for a traced caller.
+    assert "if _f_concrete(kind == _kinds_mod.I_TWICE):" in emitted, emitted
     out = tmp_path / "emitted"
     out.mkdir()
     for path, content in candidate.files.items():
@@ -1864,9 +1865,6 @@ def test_a_dispatch_on_a_module_constant_through_its_alias_is_static(tmp_path: P
         x = np.array([1.0, 3.0])
         assert np.asarray(module.pick(2, 2, x)).tolist() == [2.0, 6.0]
         assert np.asarray(module.pick(2, 1, x)).tolist() == [1.0, 3.0]
-        print(
-            "ANCHOR_SPELLING", [line for line in anchor_src.splitlines() if "I_TWICE" in line][:2]
-        )
     finally:
         sys.path.remove(str(out))
         for name in list(sys.modules):

From b811729627727551c7469b39d2f86dac2cb96af3 Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Sat, 5 Sep 2026 03:27:48 -0600
Subject: [PATCH 72/73] jax: a static test is spelled with Python logic,
 decided on its leaves

The expression mapping spells .not./.and./.or. as jnp.logical_*; under
jit those stage a constant into a tracer that no Python if can
convert (interpolation's .not. l_quintic_poly_interp took the whole
step down). The Python form of a static branch is now spelled with
Python's not/and/or, _f_concrete decides on the leaves (the
comparisons, the names), and the static predicate sees through the
mapped calls so a negated static comparison is static too.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/jax/backend.py | 52 +++++++++++++++++++++++++++--
 src/recast/transform/jax/runtime.py | 16 +++++----
 tests/test_jax_transform.py         | 19 +++++++----
 3 files changed, 71 insertions(+), 16 deletions(-)

diff --git a/src/recast/transform/jax/backend.py b/src/recast/transform/jax/backend.py
index 10625c9..f4e77df 100644
--- a/src/recast/transform/jax/backend.py
+++ b/src/recast/transform/jax/backend.py
@@ -495,6 +495,42 @@ def add(n):
     return out
 
 
+def _jnp_logic(node) -> str | None:
+    """``jnp.logical_not/and/or(...)``: which, or None."""
+    if (
+        isinstance(node, ast.Call)
+        and isinstance(node.func, ast.Attribute)
+        and isinstance(node.func.value, ast.Name)
+        and node.func.value.id == "jnp"
+        and node.func.attr in ("logical_not", "logical_and", "logical_or")
+        and not node.keywords
+    ):
+        return node.func.attr
+    return None
+
+
+def _python_logic(node):
+    """The expression mapping's ``jnp.logical_*`` back to Python's ``not``,
+    ``and``, ``or`` -- for a test a Python ``if`` evaluates at trace time,
+    where a jnp op on a constant would be staged into a tracer."""
+    which = _jnp_logic(node)
+    if which == "logical_not" and len(node.args) == 1:
+        return ast.UnaryOp(op=ast.Not(), operand=_python_logic(node.args[0]))
+    if which in ("logical_and", "logical_or") and len(node.args) == 2:
+        op = ast.And() if which == "logical_and" else ast.Or()
+        return ast.BoolOp(op=op, values=[_python_logic(a) for a in node.args])
+    return node
+
+
+def _logic_leaves(node) -> list[ast.expr]:
+    """The operands under the logical operators: what has to be concrete
+    for the Python form to be evaluable."""
+    which = _jnp_logic(node)
+    if which is not None:
+        return [leaf for a in node.args for leaf in _logic_leaves(a)]
+    return [node]
+
+
 def _module_constant(node) -> bool:
     """``_mod.IIPDF_ADG1``: a use-associated module constant through its
     module alias -- a Python value at trace time, like the bare upper-case
@@ -559,6 +595,10 @@ def constant(node):
         return all(_static_test(v, statics) for v in test.values)
     if isinstance(test, ast.UnaryOp) and isinstance(test.op, ast.Not):
         return _static_test(test.operand, statics)
+    if _jnp_logic(test) is not None:
+        # The expression mapping ran first: ``.not. ( k == I )`` arrived as
+        # ``jnp.logical_not(k == I)``. Static when its operands are.
+        return all(_static_test(a, statics) for a in test.args)
     return constant(test)
 
 
@@ -1030,7 +1070,7 @@ def lower_if(self, s, depth):
         if not body:
             body = [ast.Pass()]  # a Python if needs a body; the else is the point
         if _static_test(s.test):
-            return [ast.If(test=s.test, body=body, orelse=orelse or [])]
+            return [ast.If(test=_python_logic(s.test), body=body, orelse=orelse or [])]
         if _static_test(s.test, self.statics):
             # Over the kernel's static scalar arguments: a Python if when
             # the kernel is called through its jit wrapper (the arguments
@@ -1039,7 +1079,13 @@ def lower_if(self, s, depth):
             # takes its level indices as arguments, and its caller's loop
             # index reaches it as a tracer). Decided at trace time by the
             # runtime's ``_f_concrete``.
-            python_form: list[ast.stmt] = [ast.If(test=s.test, body=body, orelse=orelse or [])]
+            # The Python if is spelled with Python logic: ``jnp.logical_not``
+            # of a constant is staged under jit like any other op, a tracer
+            # no ``if`` can convert. Whether the leaves (the comparisons,
+            # the names) are concrete is what decides the form.
+            python_form: list[ast.stmt] = [
+                ast.If(test=_python_logic(s.test), body=body, orelse=orelse or [])
+            ]
             try:
                 cond_form = self._cond_form(s, copy.deepcopy(body), copy.deepcopy(orelse))
             except JaxQueue:
@@ -1050,7 +1096,7 @@ def lower_if(self, s, depth):
                 ast.If(
                     test=ast.Call(
                         func=ast.Name(id="_f_concrete", ctx=ast.Load()),
-                        args=[copy.deepcopy(s.test)],
+                        args=[copy.deepcopy(leaf) for leaf in _logic_leaves(s.test)],
                         keywords=[],
                     ),
                     body=python_form,
diff --git a/src/recast/transform/jax/runtime.py b/src/recast/transform/jax/runtime.py
index 6aad27b..3c05951 100644
--- a/src/recast/transform/jax/runtime.py
+++ b/src/recast/transform/jax/runtime.py
@@ -156,13 +156,15 @@ def _f_vpow(a, b):
     return jnp.asarray(a) ** b
 
 
-def _f_concrete(x):
-    """Whether a value is known at trace time -- a Python or NumPy scalar,
-    or a concrete array -- rather than a tracer. A branch over a kernel's
-    static scalar argument is a Python ``if`` when the kernel runs through
-    its jit wrapper and a ``lax.cond`` when another kernel's traced body
-    calls its implementation; this decides which."""
-    return not isinstance(x, jax.core.Tracer)
+def _f_concrete(*values):
+    """Whether every value is known at trace time -- a Python or NumPy
+    scalar, or a concrete array -- rather than a tracer. A branch over a
+    kernel's static scalar argument is a Python ``if`` when the kernel runs
+    through its jit wrapper and a ``lax.cond`` when another kernel's traced
+    body calls its implementation; this decides which, over the leaves of
+    the test (its comparisons and names), since the test itself is spelled
+    with Python logic the Python ``if`` evaluates."""
+    return not any(isinstance(x, jax.core.Tracer) for x in values)
 
 
 def _f_fori(lo, hi, body, init):
diff --git a/tests/test_jax_transform.py b/tests/test_jax_transform.py
index 28689bb..260e799 100644
--- a/tests/test_jax_transform.py
+++ b/tests/test_jax_transform.py
@@ -1803,7 +1803,7 @@ def test_a_static_branch_whose_arms_lowered_to_nothing_is_no_branch(tmp_path: Pa
   implicit none
   integer, parameter :: I_PLAIN = 1
   integer, parameter :: I_TWICE = 2
-  logical :: l_noisy = .false.
+  logical, parameter :: L_QUINTIC = .false.
 end module kinds_mod
 """
 
@@ -1825,6 +1825,9 @@ def test_a_static_branch_whose_arms_lowered_to_nothing_is_no_branch(tmp_path: Pa
       y = 0.0d0
       call complain( n, y )
     end if
+    if ( .not. ( kind == I_PLAIN ) ) then
+      y = y + 1.0d0
+    end if
   end subroutine pick
 end module dispatch_mod
 """
@@ -1849,21 +1852,25 @@ def test_a_dispatch_on_a_module_constant_through_its_alias_is_static(tmp_path: P
     facts = frontend.analyze(unit, tmp_path)
     candidate = TreeToJax(TreeConventions()).apply(unit, facts, {"root": str(tmp_path)})
     assert "pick" in candidate.notes["jax"]["kernels"], candidate.notes["jax"]
-    emitted = candidate.files[Path("dispatch_mod_jax.py")].decode()
-    # The dual form: a Python if when the kind is concrete (through the jit
-    # wrapper it is), the lax.cond only for a traced caller.
-    assert "if _f_concrete(kind == _kinds_mod.I_TWICE):" in emitted, emitted
     out = tmp_path / "emitted"
     out.mkdir()
     for path, content in candidate.files.items():
         (out / path.name).write_bytes(content)
+    emitted = candidate.files[Path("dispatch_mod_jax.py")].decode()
+    # The dual form: a Python if when the kind is concrete (through the jit
+    # wrapper it is), the lax.cond only for a traced caller.
+    assert "if _f_concrete(kind == _kinds_mod.I_TWICE):" in emitted, emitted
+    # ``.not. ( kind == I_PLAIN )``: the Python form's test is Python's
+    # ``not``, not jnp.logical_not, which jit stages into a tracer no Python
+    # if can convert (interpolation's ``.not. l_quintic_poly_interp``).
+    assert "if not kind == _kinds_mod.I_PLAIN:" in emitted, emitted
     sys.path.insert(0, str(out))
     try:
         module = importlib.import_module("dispatch_mod_jax")
         import numpy as np
 
         x = np.array([1.0, 3.0])
-        assert np.asarray(module.pick(2, 2, x)).tolist() == [2.0, 6.0]
+        assert np.asarray(module.pick(2, 2, x)).tolist() == [3.0, 7.0]
         assert np.asarray(module.pick(2, 1, x)).tolist() == [1.0, 3.0]
     finally:
         sys.path.remove(str(out))

From 3573059ba9df6823210bd86aac5aff17d6f9ec36 Mon Sep 17 00:00:00 2001
From: lewisychen <yueqichen.0x0@gmail.com>
Date: Sat, 5 Sep 2026 03:55:37 -0600
Subject: [PATCH 73/73] jax: a while loop's exit flag starts at the function
 top; its carries are its lowered stores

A DO WHILE under an IF inside a DO (interpolation's grid search): the
exit flag is a carry of the enclosing cond and needs a value before the
branch, at the top of the function like the goto-region flags, with the
one beside the loop resetting it on entry. And the loop's carries are
what its lowered body stores: a subscript store is a store to its base
only once lowered, and an inner loop's break flag and kept index are
the body's own.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Cne4hrGgYqhTMzVgzJtXYd
---
 src/recast/transform/jax/tree.py | 22 ++++++++---
 tests/test_jax_transform.py      | 67 ++++++++++++++++++++++++++++++++
 2 files changed, 83 insertions(+), 6 deletions(-)

diff --git a/src/recast/transform/jax/tree.py b/src/recast/transform/jax/tree.py
index 076037e..fc294ff 100644
--- a/src/recast/transform/jax/tree.py
+++ b/src/recast/transform/jax/tree.py
@@ -1717,9 +1717,11 @@ def _rewritten_body(fn: ast.FunctionDef, rewrite: _Rewrite) -> list[ast.stmt]:
         n for n, kind in (rewrite.inits or {}).items() if kind == "int32"
     )
     lowered = _WhileLoops(_integer_inits(fn), int_locals).visit_block(lowered)
-    # The goto-region flags are synthetic locals with no UB-guard init; an
-    # enclosing loop carries them, and the initial carry tuple needs a value
-    # before the loop. Every flag starts False at the top of the function.
+    # The goto-region and while-exit flags are synthetic locals with no
+    # UB-guard init; an enclosing loop or branch carries them (a while inside
+    # an if inside a do: CLUBB's grid interpolation), and the initial carry
+    # tuple needs a value before it. Every flag starts False at the top of
+    # the function; the one beside its loop resets it on every entry.
     flags = sorted(
         {
             n.id
@@ -1727,7 +1729,7 @@ def _rewritten_body(fn: ast.FunctionDef, rewrite: _Rewrite) -> list[ast.stmt]:
             for n in ast.walk(stmt)
             if isinstance(n, ast.Name)
             and isinstance(n.ctx, ast.Store)
-            and re.fullmatch(r"_(?:skip|restart)_\d+", n.id)
+            and re.fullmatch(r"_(?:skip|restart|done)_\d+", n.id)
         }
     )
     lowered = [
@@ -2649,12 +2651,20 @@ def visit_While(self, node: ast.While) -> ast.AST:
                 exited = True
         for statement in guarded:
             ast.fix_missing_locations(statement)
-        carried = [n for n in _assigned_names(guarded) if n != done]  # type: ignore[no-untyped-call]
-        state = [*carried, done]
         try:
             lowered = KernelLowerer().lower_block(guarded, 1)  # type: ignore[no-untyped-call]
         except JaxQueue as why:
             raise NotFlat(f"while loop body: {why}") from why
+        # The carries are what the *lowered* body stores: a subscript store
+        # (``idx(i) = k`` in a search loop) is a store to its base only once
+        # lowered. An inner loop's break flag and kept index are the body's
+        # own, initialized beside that loop, never a carry of this one.
+        carried = [
+            n
+            for n in _assigned_names(lowered)  # type: ignore[no-untyped-call]
+            if n != done and not re.fullmatch(r"_(?:brk|kx)_\d+", n)
+        ]
+        state = [*carried, done]
         unpack = ast.Assign(
             targets=[
                 ast.Tuple(elts=[ast.Name(id=n, ctx=ast.Store()) for n in state], ctx=ast.Store())
diff --git a/tests/test_jax_transform.py b/tests/test_jax_transform.py
index 260e799..c5af316 100644
--- a/tests/test_jax_transform.py
+++ b/tests/test_jax_transform.py
@@ -1879,6 +1879,73 @@ def test_a_dispatch_on_a_module_constant_through_its_alias_is_static(tmp_path: P
                 sys.modules.pop(name, None)
 
 
+NESTED_WHILE = """\
+module nestedwhile_mod
+  implicit none
+contains
+  subroutine locate( n, m, x, grid, idx )
+    integer, intent(in) :: n, m
+    real(8), intent(in) :: x(n), grid(m)
+    integer, intent(out) :: idx(n)
+    integer :: i, k
+    logical :: calc_done
+    idx = 0
+    do i = 1, n
+      if ( x(i) > 0.0d0 ) then
+        k = 1
+        calc_done = .false.
+        do while ( .not. calc_done .and. k <= m )
+          if ( grid(k) >= x(i) ) then
+            idx(i) = k
+            calc_done = .true.
+          end if
+          k = k + 1
+        end do
+      end if
+    end do
+  end subroutine locate
+end module nestedwhile_mod
+"""
+
+
+def test_a_while_inside_a_branch_inside_a_loop_has_its_flag_before_the_branch(
+    tmp_path: Path,
+) -> None:
+    """interpolation's lin_interp_between_grids: a DO WHILE search under an
+    IF inside a DO. The while's exit flag is a carry of the enclosing cond,
+    so it needs a value before the branch -- at the top of the function,
+    like the goto-region flags -- not only beside its loop."""
+    import importlib
+    import sys
+
+    from recast.transform.jax.tree import TreeToJax
+    from recast.transform.numpy.tree import TreeConventions
+
+    (tmp_path / "nestedwhile_mod.f90").write_text(NESTED_WHILE)
+    frontend = FortranFrontend(flatten=True)
+    unit = next(u for u in frontend.discover(tmp_path) if u.uid == "fortran:nestedwhile_mod")
+    facts = frontend.analyze(unit, tmp_path)
+    candidate = TreeToJax(TreeConventions()).apply(unit, facts, {"root": str(tmp_path)})
+    assert candidate.notes["jax"]["kernels"] == ["locate"], candidate.notes["jax"]
+    out = tmp_path / "emitted"
+    out.mkdir()
+    for path, content in candidate.files.items():
+        (out / path.name).write_bytes(content)
+    sys.path.insert(0, str(out))
+    try:
+        module = importlib.import_module("nestedwhile_mod_jax")
+        import numpy as np
+
+        x = np.array([0.5, -1.0, 2.5, 9.0])
+        grid = np.array([1.0, 2.0, 3.0])
+        got = np.asarray(module.locate(4, 3, x, grid)).tolist()
+        assert got == [1, 0, 3, 0]
+    finally:
+        sys.path.remove(str(out))
+        for suffix in ("_jax", "_numpy", "_jax_runtime", "_constants"):
+            sys.modules.pop(f"nestedwhile_mod{suffix}", None)
+
+
 SILENT_CHECK = """\
 module silent_mod
   implicit none