fix(core): a step that changes the layout of the state is refused (MADD-ANO-220)

[NicholasEhsanRoy]

Resolves MADD-ANO-220: GraphManager.step, run and run_adaptive stored a stepped state whose layout differed from the state it replaced (shipped 0.1.0 to 0.3.1, silent). They now raise ValueError naming the node and the leaf, and store nothing.

Decision needed before merge (PR left as draft for it)

A node whose update returns a field its initial_state() does not build is now refused at its first step. The guard compares keys as well as shapes (as the saved patch did), so a state that gains a field is a layout change. One existing test pinned the opposite as working: tests/cloud/multigpu/test_unstructured_layout_contract.py::test_a_node_without_an_initial_integral_steps_in_a_graph (verification of MADD-ANO-040), where a node emits a domain integral that initial_state() lacks and gm.step() stored it. docs/developer_guide/sharding_topology.md described that as "a graph's second step". No stock node and no shipped example does this (domain_integral_fields is used by none of them; the sweeps below found no other graph).

What this branch does: the test keeps what it protected (the wrapper is handed its own output back, three jitted update calls, same integral) and now also asserts the graph's refusal; the doc says a node in a graph gives its integral an initial value; the declared form was already tested in graphs and scans (test_a_node_declaring_its_integral_runs_in_a_graph_like_the_unsharded_node).

The alternative, if emitting an undeclared field through step() is to stay supported: do not report appears only after the update in GraphManager._refuse_layout_drift (a few lines), and revert that test and doc sentence. The cost: such a state still does not scan, and its checkpoint has a field the reset graph lacks, which is the defect class of 220.

What the saved patch needed

p4_34_step_guard.patch applied cleanly with git apply --3way. Four things were changed after reading and probing it; each has a regression test and a mutant:

1. Once per trace, not once per compile. The patch's flag was cleared by the first step after a compile. An external input of another shape at a later step retraces the step and broadcasts a scalar leaf, and that was still stored (reproduced: add_external_input("a", "drive"), two good steps, then drive=jnp.ones(3)). The comparison is now due whenever (compile generation, trace count) differs from the last one compared. Cost on every other step: one tuple comparison. params= of another shape was already refused by the params shape check.
2. Kinds of dtype are compared (the patch passed dtypes=False). An integer leaf returned as a float was stored the same way, is refused by the scans, and its checkpoint is refused after a reset ("holds 0.5 (float32), which this graph's int32 cannot hold"): same class.
3. run_adaptive compares the first step it keeps, before its callback, observers and the dt_min warning. The patch compared at the end of the run, after every callback and EVENT_STEP observer had been handed the reshaped state.
4. A PRNG key in the state. _leaf_layout raised on an extended dtype (key<fry>), so the guard refused every step of a graph that carries a key (found by the sweep: test_step_program_digests.py, chain-into-ring/choice-0). A key is now compared as a kind of its own (same implementation). POST /graph/nodes had been refusing such a node through that internal error; it now refuses it by name (a key has no JSON form, so no state reply could carry it; a graph with a key handed to the server in process answers 500 on /sim/step on the base too, which is a separate, existing matter).

The layout rule

Keys (a missing or a new leaf), shapes, and the kind of dtype (boolean, integer, float, complex, or a PRNG key of one implementation). Width is not compared: under x64 a stock node returns float64 for a float32 state and a width check would refuse them all. A width change is still stored silently by step (pinned: test_a_dtype_of_another_width_is_not_what_the_comparison_refuses); that is the x64 carry matter of MADD-ANO-017 and 190, not this entry.

Host-side only: shapes and dtypes are read from the returned state (tracers under a transform included); nothing is traced, no node's update is called.

Entry points

Entry point	Covered by	Test
step, run	_store_stepped_state, once per trace	test_the_step_entry_points_keep_the_shape_of_every_state_leaf[step, run]
run_adaptive	first kept step of each call	same test, [run_adaptive]
retrace at a later step	trace count in the stamp	test_a_later_step_that_is_traced_again_is_compared_again
recompile	compile generation in the stamp	test_the_step_of_a_graph_compiled_again_is_compared_again
under jax.grad	same path, on tracers	test_a_step_inside_a_transform_is_compared_on_the_traced_state
coupling group	same path; group bookkeeping keys are leaves	test_a_graph_with_a_coupling_group_is_held_to_its_layout_too
sharded node	same path	tests/cloud/multigpu/test_a_sharded_step_keeps_the_state_layout.py
REST /sim/step, /sim/run	reach gm.step / gm.run; 400, nothing stored	test_a_node_the_door_never_saw_is_refused_where_the_graph_steps
FMU sidecar and bridge	do not reach gm.step (they run gm._compiled_step on their own state): own comparison in FmuSidecar._advanced	tests/fmi/test_a_sidecar_step_keeps_the_state_layout.py
run_scan, run_scan_with_history, run_adaptive_scan	unchanged, TypeError (carry)	test_the_scan_entry_points_refuse_a_step_that_changes_a_leafs_shape

The sidecar was a sibling site the entry did not name: build_model_description accepted the broadcasting ball and sidecar.step stored position as (2,). It compares once per trace for a graph's current compiled step, and at every step for any other callable or for the step of a graph compiled again since.

On refusal the graph holds the same state object as before (gm._state is held; the clock and step count of a multi-rate or coupled graph live in it), no observer is told of a step, no callback runs, and the refusal repeats until the node is fixed. After an external input of the wrong shape is refused, stepping on with the right one matches a fresh graph to the bits.

How the REST door relates: _dry_run_node stays. It refuses before a node is added, traces abstractly, and names the parameter to blame. The graph's guard is the backstop for every graph that did not pass that door (built in process, loaded, edited in process).

The compile-count fixture

AvalChangesOnce grew x from (1,) to (2,). It now starts x as a Python float: the first step is traced against a weakly typed float32 scalar and returns a strongly typed one, so the second step retraces once and the count settles at 2. Shape and kind never change; the state saves, reloads and scans. Measured: trace_count == 2, compile_counts(gm).retrace_count == 2, and the child-process gate tests built on the fixture pass unchanged.

Graphs the guard refused

Sweeps with a plugin that logs every refusal, swallowed or not (jax 0.11.0, 3 cores, per-push lane): tests/core (four chunks, 8,365 passed), tests/nodes (877), tests/test_examples_smoke.py (273), tests/fmi (838), tests/adaptive, tests/verification and tests/cloud/multigpu without the run_pod files (1,344 together), and the tests/property, tests/surrogates, tests/viz, tests/usd files that never name /cloud (1,412). The tests/api files that never name /cloud ran in one long process that was killed at the 10 GB cap 73% of the way through a combined run; nothing was refused before that, and CI is the check for the rest.

Graph	What it was	What was done
AvalChangesOnce (test_compile_counts.py)	reshaped on purpose to retrace	replaced as above
chain-into-ring/choice-0 (test_step_program_digests.py)	a PRNG key in the state; the comparison raised on it	comparison fixed (item 4)
_Source(8) under ShardedUnstructuredNode	emits an integral initial_state() lacks	the decision above

No stock node and no shipped example was refused.

Programs, counts, bit identity

• scripts/capture_step_programs.py --check, fresh process each: 24 graphs unchanged on jax 0.10.2, 0.11.0 and 0.11.2 (before and after merging the base).
• scripts/compile_counts.py --check on jax 0.10.2: counts match, 6 workloads.
• test_the_comparison_is_made_once_per_trace_and_calls_no_update: nine steps, one trace, one update call, one comparison.
• The sidecar's five steps equal the graph's five to the bits; the existing bit-identity tests in the swept directories pass.

Mutants: 17 seeded, 17 caught

Comparison skipped; stamp without the compile generation (the "never reset on recompile" fault); stamp without the trace count; refused but stored; a refusal that uses the comparison up; kind of dtype not compared; run_adaptive not compared; run_adaptive compared only after its callbacks; compared at every step; sidecar: skipped, trusting a non-graph step, ignoring a later compile, comparing at every graph step, refusal uses the comparison up; a key not a kind of its own; any two keys one kind; the REST door taking a keyed state.

Code paths newly enabled or changed

• A graph with a PRNG key in its state reaches the layout comparison without an internal error (it steps as on the base).
• Released behaviour changed, to refuse a silently wrong result (CHANGELOG ### Changed): a step that changes a leaf's shape or kind of dtype, or the state's fields, raises where it was stored; FmuSidecar.step and a bridge step likewise (error reply, nothing committed).
• One existing CHANGELOG line (the REST door of 220) said the defect "stays open in process"; it was edited in place because it is no longer true.

Sibling sites checked

GraphManager: every _store_state caller (step, run, run_adaptive: guarded; run_scan, run_scan_with_history, run_adaptive_scan: loud already; run_sweep stores nothing). profiler.compile_counts and sim_profile step through gm.step. FmuSidecar.step and FmuTcpBridge (all advancing goes through _advanced): guarded here. examples/basics/bouncing_ball.py calls _compiled_step on a local state and stores nothing in the graph.

Not run locally

The full suite, the slow lane beyond the files that cite the changed symbols (test_compile_counts.py, test_bridge_master_dt_is_the_graph_step.py, test_differential_fmu.py: their slow tests ran once, all passed), and the test files that send requests to /cloud. Pyright 1.1.414 on the two changed source files: 0 errors (local numpy, not CI's).

One observation, not acted on: a graph with a sharded node beside a plain one traces its step twice (the plain node's leaves come back placed on the mesh), so it compiles twice per run.

🤖 Generated with Claude Code

https://claude.ai/code/session_013UkCde7g23gTziUjYAvnKD