icon_sc.core.functional — the functional (JAX) lowering

The §8.5-8.6 functional lowering (F-tier), proven at column scale (SPEC S10).

Layout (architecture §8.5; jax-touching core code lives here and nowhere else in icon_sc.core — importing icon_sc.core itself never imports jax):

  • icon_sc.core.functional.pytree — generated frozen-dataclass PyTrees: StateTree from the state/vault schema plus explicit carry, ParamTree from the components’ params declarations.

  • icon_sc.core.functional.compile — composition → pure step_fn(StateTree, ParamTree, StaticArgs) -> StateTree and the scan_window (lax.scan + per-step jax.checkpoint) window builder.

  • icon_sc.core.functional.rules — implicit-function-theorem helpers for fixed points (lax.custom_root wrappers; the §8.6 custom route).

fp64 is the default for gradient work (§8.6): the compile entry point warns when jax runs in fp32 (jax.config.update("jax_enable_x64", True)).

Functional compile: composition → pure step_fn + scan_window (§8.5, SPEC S10).

functional_compile(composition, state, timestep=...) is the F-tier consumer of the composition walk (the same visit(plan_builder) double-dispatch the S05 plan compiler uses): instead of an imperative op list it emits a pure JAX function over explicit PyTrees,

step_fn: (StateTree, ParamTree, StaticArgs) -> StateTree # one Δt, traced window = scan_window(step_fn, n_steps, remat=”per_step”) # lax.scan + checkpoint

The semantic mapping is mechanical (§8.5): sequential updates become functional updates (memory recovered by donate_argnums under jit); cadence wrappers become carry (cached output + last-fire phase) selected by jnp.where on a carried step counter — the trace is static across step signatures; monitors and time are outside the trace.

Component protocol (§8.6 native/custom): a compilable component provides a pure functional_call(inputs, params, dt=...) -> outputs — inputs keyed by its contract names, outputs the flat union of its output dicts — co-located with its imperative kernel and drawing on the same scheme-constants module, plus functional_params() -> {name: default} for its params declarations. Components whose contracts declare no native/custom output are none under the differentiability contract: per differentiated region the composition-time policy is "error" (default) or "stop_gradient" — explicit, warned, and stamped into provenance. Gradient truncation is never silent.

exception icon_sc.core.functional.compile.FunctionalCompileError

The composition cannot be lowered to a pure function (names the node).

class icon_sc.core.functional.compile.FunctionalProgram(step_fn, state, params, static, state_type, param_type, carry_names, provenance)

The compiled F-tier program (frozen interface, SPEC S10).

step_fn(state, params, static) -> state is pure and traced; state / params are instances of the generated state_type / param_type PyTrees; provenance stamps every composition-time decision (cadence folds, stop_gradient truncations, freezes) — never silent (§8.6).

Parameters:
class icon_sc.core.functional.compile.StaticArgs(dt)

Per-call static/scalar arguments of the pure step (frozen interface, SPEC S10).

dt must equal the timestep the program was compiled against: cadence periods are folded into the trace as step counts at compile time.

Parameters:

dt (float)

icon_sc.core.functional.compile.functional_compile(composition, state, *, timestep, policy='error')

Lower a composition to a pure step_fn over explicit PyTrees (SPEC S10).

composition is one composition node or an ordered sequence of nodes (the §5.1 loop-body order — e.g. the SCM preset’s (slow, core, fast)); every node lowers through the same visit(plan_builder) protocol the S05 plan compiler consumes. state is the schema-representative boundary state the trees are generated from. policy governs differentiable: 'none' components per §8.6: "error" (default) or "stop_gradient".

Parameters:
Return type:

FunctionalProgram

icon_sc.core.functional.compile.scan_window(step_fn, n_steps, *, remat='per_step', ys_of=None)

A multi-step window over step_fn: lax.scan + checkpoint policy (§8.5).

Returns window(state, params, static). remat="per_step" wraps the step in jax.checkpoint (reverse-mode memory ≈ one extra forward, §8.7); remat=None stores all activations. ys_of optionally maps each post-step state to a per-step observable; the window then returns (final_state, stacked_ys) instead of final_state.

Parameters:
Return type:

Callable[[…], Any]

Generated frozen-dataclass PyTrees: StateTree and ParamTree (§8.5, SPEC S10).

StateTree is derived from the state/vault schema — one leaf per slot — extended with explicit carry: every leaf a component declared in functional_state() is surfaced into the tree (tension T7: §4.5’s privacy is an imperative-tier convenience, demoted here by contract). ParamTree holds the tunable scheme constants of the §8.6 params declarations, distinct from state — calibration never smuggles constants through state fields.

Leaf ordering is deterministic: leaves are sorted by canonical name at type generation (PLAN pitfall — carry ordering must not depend on dict insertion order). Canonical names (icon:qnc, fcarry/0/...) are not identifiers; they are sanitized into attribute names with the canonical → attribute map kept on the generated class (__icon_sc_leaves__).

icon_sc.core.functional.pytree.build_param_tree(params, *, type_name='ParamTree')

ParamTree type + instance from a flat name → default-value mapping (§8.6).

Parameters:
Return type:

tuple[type, Any]

icon_sc.core.functional.pytree.build_state_tree(state, extra_leaves=None, *, type_name='StateTree')

StateTree type + instance from a state dict (one leaf per slot) plus carry.

state is a boundary state (dict of DataArrays; time is skipped — the F-tier trace has no clock, cadence rides in the carry). extra_leaves are the explicit-carry leaves and any compiler-seeded slots (already arrays).

Parameters:
Return type:

tuple[type, Any]

icon_sc.core.functional.pytree.make_pytree_type(type_name, leaf_names)

Generate a frozen-dataclass PyTree type with one field per leaf name.

Leaves are sorted by canonical name; the generated class is registered with jax.tree_util.register_dataclass (all fields are data fields) and carries __icon_sc_leaves__: the (canonical_name, attribute_name) pairs in field order.

Parameters:
Return type:

type

icon_sc.core.functional.pytree.mapping_of(tree)

The canonical-name → leaf mapping of a generated PyTree instance.

Parameters:

tree (Any)

Return type:

dict[str, Any]

icon_sc.core.functional.pytree.sanitize_leaf_name(name)

A valid Python attribute name for one canonical leaf name.

Parameters:

name (str)

Return type:

str

icon_sc.core.functional.pytree.tree_of(cls, values)

Instantiate a generated PyTree type from a canonical-name → leaf mapping.

Parameters:
Return type:

Any

Implicit-function-theorem rules for fixed points (§8.6 custom route, SPEC S10).

Implicit structure gets IFT treatment rather than unrolling: the fixed point differentiates through a lax.custom_root-style rule, not through recorded Newton iterations — which also sidesteps while_loop’s reverse-mode prohibition (the solver runs inside custom_root’s non-differentiated primal).

Verified on the pinned jax (REFERENCES.lock jax): lax.custom_root with a linear elementwise tangent_solve supports both jvp and vjp (reverse mode derives from the JVP by transposition), with a bounded lax.while_loop inside solve.

icon_sc.core.functional.rules.implicit_fixed_point(residual, x0, solve)

Solve residual(x) == 0 elementwise; differentiate via the IFT (§8.6).

residual must be elementwise in x (its Jacobian at the solution is diagonal): the linearized-residual solve is then y / g(1), which is linear — exactly what lax.custom_root needs to derive both the JVP and, by transposition, the VJP. Everything residual closes over is differentiated through the implicit function; solve is primal-only and may use lax.while_loop.

Parameters:
Return type:

Any

icon_sc.core.functional.rules.masked_newton_solve(residual, residual_prime, x0, *, active0, tolerance, max_iter)

Elementwise Newton with per-point convergence freeze (primal-only solver).

Mirrors the granule-style masked iteration (icon4py satad, REFERENCES.lock icon4py-satad-stencils): points update only while their mask is active; the mask deactivates when |Δx| <= tolerance; the global loop runs while any point is active, bounded by max_iter (a data-dependent raise is not expressible under jit — the T0 granule raises ConvergenceError instead; deviation recorded in STATUS S10). Meant to be passed as the solve of implicit_fixed_point(), so its iterations are never differentiated.

Parameters:
Return type:

Any