From e5bda8e554e15e9122dc11cb0114cdf80b7df948 Mon Sep 17 00:00:00 2001 From: Paul Fidika Date: Thu, 6 Aug 2026 19:59:41 -0600 Subject: [PATCH] =?UTF-8?q?pgw#993:=20one=20flattening=20rule=20=E2=80=94?= =?UTF-8?q?=20`carried=5Fby`=20resolves=20through=20the=20same=20expansion?= =?UTF-8?q?=20`dynamic=5Fshapes`=20mirrors?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit MEASURED (gen-worker 0.93.2, pod `7evvazd2xplzml`, A100-SXM4-80GB, $0.4655): the z-image AOT mint passed every earlier gate, entered the AOT recipe, and refused in `trace_graph` — exit=2, `deterministic`, four identical attempts: entry 'transformer/adapter=true,cfg=true': declared-range gate: declared dynamic dim names input 'x', which is not a user input of the exported program (inputs: ['cap_feats_0', 'cap_feats_1', 'lora_a', 'lora_b', 't', 'x_0', 'x_1']) Two SDK features that each work and could not compose. `Input.repeat` containers (pgw#853) are FLATTENED by torch.export into one positional user input per element, suffixed `_0`, `_1`, … for EVERY arity — a one-element container is `x_0`, never `x` (re-measured here on a real export, both arms). `Dim.carried_by` names its input by the DECLARED name, and the declared-range gate resolved that name against the exported program. That makes a `Dim` carried by a repeated container unsatisfiable BY CONSTRUCTION: naming `x_0` breaks the `cfg=false` arm, and dropping the dims forfeits `dynamic-collapse`, which is the reason z-image is 2 entries and not 20. The fix is the invariant, not the patch. `aot_mint.exported_input_names` is now the ONE expansion rule: `dynamic_shapes_spec` mirrors the container structure with it, and `declared_range_gaps` + `lifted_input_gaps` resolve declared names against exported ones with it — all three from the same per-arm arity map (`aot_declaration.container_arities`) the example feed was built from, computed once at the mint call site. A gate that resolves declared names against exported names must share the name-mangling with the code that produced them; two independent spellings of one mapping is the defect class, which is why the lifted-input gate is in this change rather than left for the next rented pod. Every element is gated, not merely found: a per-element specialization is still refused, under the element's own name (`x_1[2] exported as the STATIC value 64 …`). Non-container declarations resolve exactly as before (pgw#846). RED, off-GPU, on the real path (declaration -> export -> gate), reproduced on BOTH `v0.93.2` and `origin/master` and green here: N=1 user_inputs=['x_0', 't'] gate(declared names) -> refusal N=2 user_inputs=['x_0', 'x_1', 't'] gate(declared names) -> refusal N=1/N=2 with the arity map -> [] tests/test_dim_flattening_pgw993.py drives the pod's program as a double (user inputs verbatim, refusal string verbatim) and then exports for real on CPU for both arms of the fork. --- changelog.d/pgw993.md | 24 ++ src/gen_worker/aot_mint.py | 245 ++++++++++++-------- tests/test_dim_flattening_pgw993.py | 335 ++++++++++++++++++++++++++++ 3 files changed, 507 insertions(+), 97 deletions(-) create mode 100644 changelog.d/pgw993.md create mode 100644 tests/test_dim_flattening_pgw993.py diff --git a/changelog.d/pgw993.md b/changelog.d/pgw993.md new file mode 100644 index 00000000..98494d10 --- /dev/null +++ b/changelog.d/pgw993.md @@ -0,0 +1,24 @@ +- **pgw#993: a `Dim` carried by a `repeat=` container was unsatisfiable by + construction, so z-image could never mint an AOT cell.** Measured on + gen-worker 0.93.2, pod `7evvazd2xplzml` (A100-SXM4-80GB, \$0.4655): the mint + passed every earlier gate, entered the AOT recipe and refused in + `trace_graph`, exit=2, `deterministic`, four identical attempts — + `entry 'transformer/adapter=true,cfg=true': declared-range gate: declared + dynamic dim names input 'x', which is not a user input of the exported + program (inputs: ['cap_feats_0', 'cap_feats_1', 'lora_a', 'lora_b', 't', + 'x_0', 'x_1'])`. Two features that each worked: `Input.repeat` containers + (pgw#853) export as one flattened user input per element (`x_0`, `x_1`, …, + suffixed for EVERY arity — a one-element container is `x_0`, never `x`), + while `Dim.carried_by` names its input by the DECLARED name and the + declared-range gate resolved that name against the exported program. No + declaration edit could satisfy both arms of a CFG fork, and dropping the + dims would have forfeited `dynamic-collapse` — the reason z-image is 2 + entries and not 20. `aot_mint.exported_input_names` is now the ONE + flattening rule: `dynamic_shapes_spec` mirrors the container structure with + it, and `declared_range_gaps` and `lifted_input_gaps` resolve declared names + against exported ones with it, from the same per-arm arity map + (`aot_declaration.container_arities`) the example feed was built from. Every + element is gated, so a per-element specialization is still refused, by the + element's own name. The lifted-input gate is in the same change rather than + left for the next rented pod, because two independent spellings of one name + mapping is the defect class, not the symptom. diff --git a/src/gen_worker/aot_mint.py b/src/gen_worker/aot_mint.py index 6287f7c0..6a6bc2a5 100644 --- a/src/gen_worker/aot_mint.py +++ b/src/gen_worker/aot_mint.py @@ -238,6 +238,35 @@ def lifted_torch_gap(spec: ExportSpec) -> str: # --------------------------------------------------------------------------- +def exported_input_names( + name: str, containers: Optional[Mapping[str, int]] = None, +) -> Tuple[str, ...]: + """The exported program's user-input name(s) a DECLARED input name means. + + THE ONE FLATTENING RULE (pgw#993). ``torch.export`` flattens a container + argument into one positional user input per element, suffixed ``_0``, + ``_1``, … — measured for every arity, ``N=1`` included, so a one-element + container is ``x_0`` and never ``x``. A declared name therefore survives + into the exported program only when it is NOT a container. + + Every consumer that maps declared names onto exported ones goes through + here: :func:`dynamic_shapes_spec` (which mirrors the structure), + :func:`declared_range_gaps` and :func:`lifted_input_gaps` (which resolve + names against the program). Two independent spellings of one mapping is + the defect class this replaces — z-image's ``Dim(carried_by=(("x", 2),))`` + was unsatisfiable by construction, refusing every mint with "declared + dynamic dim names input 'x', which is not a user input of the exported + program (inputs: [… 'x_0', 'x_1'])" while the export it gated succeeded. + + The arity comes from the SAME class row the example feed was built from + (``aot_declaration.container_arities``), never re-guessed. + """ + arity = (containers or {}).get(name) + if arity is None: + return (str(name),) + return tuple(f"{name}_{index}" for index in range(int(arity))) + + def dynamic_shapes_spec( dims: Sequence[DynamicDim], input_names: Sequence[str], containers: Optional[Mapping[str, int]] = None, @@ -311,18 +340,16 @@ def dynamic_shapes_spec( # fine, so this was an SDK gap, not a torch limitation. The arity comes # from the SAME class row the feed was built from # (`aot_declaration.container_arities`), never re-guessed here. - arities = dict(containers or {}) out: Dict[str, Any] = {} for name in input_names: spec = by_input.get(name) - arity = arities.get(name) - if arity is None: - out[name] = spec - else: - # One entry per element; every element of a declared container - # shares the container's declared axes, which is what makes the - # elements one graph class rather than N. - out[name] = [spec for _ in range(int(arity))] + # One entry per EXPORTED element; every element of a declared + # container shares the container's declared axes, which is what makes + # the elements one graph class rather than N. The element count comes + # from `exported_input_names` — the same rule the declared-range and + # lifted-input gates resolve names with (pgw#993). + exported = exported_input_names(name, containers) + out[name] = spec if exported == (name,) else [spec for _ in exported] return out @@ -441,6 +468,7 @@ def _shape_env(program: Any) -> Any: def declared_range_gaps( program: Any, dims: Sequence[DynamicDim], + containers: Optional[Mapping[str, int]] = None, ) -> List[str]: """Named reasons the export did not honour the declared dynamic contract. @@ -466,6 +494,11 @@ def declared_range_gaps( shares a factor records nothing. This is the check the presence-only gate lacked, and it is evidence-based rather than arithmetic — see :func:`_pinning_guards` for why the arithmetic version was wrong. + + A dim carried by a ``repeat=`` container is checked on EVERY exported + element (pgw#993): the declared name is resolved through + :func:`exported_input_names`, the same flattening rule + :func:`dynamic_shapes_spec` mirrors the structure with. """ gaps: List[str] = [] shapes = _placeholder_shapes(program) @@ -474,90 +507,97 @@ def declared_range_gaps( for d in dims: if d.min == d.max: continue - shape = shapes.get(d.input_name) - if shape is None: - gaps.append( - f"declared dynamic dim names input {d.input_name!r}, which is " - f"not a user input of the exported program " - f"(inputs: {sorted(shapes)!r})") - continue - if d.axis >= len(shape): - gaps.append( - f"{d.input_name}[{d.axis}] is out of range for the exported " - f"shape {tuple(str(x) for x in shape)!r}") - continue - dim = shape[d.axis] - text = str(dim) - if text.lstrip("-").isdigit(): - gaps.append( - f"{d.input_name}[{d.axis}] exported as the STATIC value {text} " - f"but is declared dynamic [{d.min}, {d.max}] — export " - f"specialized a dim the declaration advertises as dynamic") - continue - syms = _free_symbols(dim) - declared_symbols.extend(syms) - covered = False - # The SOLVED range of the axis's full expression, when the program - # records one. This is what makes a UNIFIED relational axis (#739 / - # ie#566 §5) gate-able: wan ti2v's per-token dim solves to - # ``31*s25*s56`` with its own composite range entry, while the - # per-symbol path below would compare a governing symbol's [20, 40] - # against the declared [12400, 49600] and refuse a sound artifact. - # Composite entries are in the axis's OWN units, so the declared - # bounds compare directly, with no multiple-of scaling. - expr = getattr(getattr(dim, "node", None), "expr", None) - interval = ranges.get(expr) if expr is not None else None - if interval is not None: - try: - lo, hi = int(interval.lower), int(interval.upper) - except (TypeError, ValueError, OverflowError): - lo = hi = -1 - if lo >= 0: + for input_name in exported_input_names(d.input_name, containers): + shape = shapes.get(input_name) + if shape is None: + declared = ( + repr(d.input_name) if input_name == d.input_name + else f"{d.input_name!r} (flattened element " + f"{input_name!r})") + gaps.append( + f"declared dynamic dim names input {declared}, which is " + f"not a user input of the exported program " + f"(inputs: {sorted(shapes)!r})") + continue + if d.axis >= len(shape): + gaps.append( + f"{input_name}[{d.axis}] is out of range for the exported " + f"shape {tuple(str(x) for x in shape)!r}") + continue + dim = shape[d.axis] + text = str(dim) + if text.lstrip("-").isdigit(): + gaps.append( + f"{input_name}[{d.axis}] exported as the STATIC value " + f"{text} but is declared dynamic [{d.min}, {d.max}] — " + f"export specialized a dim the declaration advertises as " + f"dynamic") + continue + syms = _free_symbols(dim) + declared_symbols.extend(syms) + covered = False + # The SOLVED range of the axis's full expression, when the program + # records one. This is what makes a UNIFIED relational axis (#739 / + # ie#566 §5) gate-able: wan ti2v's per-token dim solves to + # ``31*s25*s56`` with its own composite range entry, while the + # per-symbol path below would compare a governing symbol's [20, 40] + # against the declared [12400, 49600] and refuse a sound artifact. + # Composite entries are in the axis's OWN units, so the declared + # bounds compare directly, with no multiple-of scaling. + expr = getattr(getattr(dim, "node", None), "expr", None) + interval = ranges.get(expr) if expr is not None else None + if interval is not None: + try: + lo, hi = int(interval.lower), int(interval.upper) + except (TypeError, ValueError, OverflowError): + lo = hi = -1 + if lo >= 0: + if lo == hi: + gaps.append( + f"{input_name}[{d.axis}] ({expr}) solved to the " + f"single value {lo} — the declared range " + f"[{d.min}, {d.max}] is advertised but the " + f"artifact admits ONE shape") + elif lo > d.min or hi < d.max: + gaps.append( + f"{input_name}[{d.axis}] ({expr}) solved to " + f"[{lo}, {hi}] which does not cover the declared " + f"[{d.min}, {d.max}] — the artifact admits less " + f"traffic than it advertises") + continue + for sym in syms: + interval = ranges.get(sym) + if interval is None: + continue + try: + lo, hi = int(interval.lower), int(interval.upper) + except (TypeError, ValueError, OverflowError): + continue if lo == hi: gaps.append( - f"{d.input_name}[{d.axis}] ({expr}) solved to the " + f"{input_name}[{d.axis}] symbol {sym} solved to the " f"single value {lo} — the declared range " f"[{d.min}, {d.max}] is advertised but the artifact " f"admits ONE shape") - elif lo > d.min or hi < d.max: + covered = True + break + # The symbol may carry a multiple-of factor (8*s95), so compare + # the DECLARED bounds against the symbol's own solved bounds + # scaled by the factor the declaration states. + factor = max(1, int(d.multiple_of or 1)) + want_lo, want_hi = d.min // factor, d.max // factor + if lo > want_lo or hi < want_hi: gaps.append( - f"{d.input_name}[{d.axis}] ({expr}) solved to " - f"[{lo}, {hi}] which does not cover the declared " - f"[{d.min}, {d.max}] — the artifact admits less " - f"traffic than it advertises") - continue - for sym in syms: - interval = ranges.get(sym) - if interval is None: - continue - try: - lo, hi = int(interval.lower), int(interval.upper) - except (TypeError, ValueError, OverflowError): - continue - if lo == hi: - gaps.append( - f"{d.input_name}[{d.axis}] symbol {sym} solved to the " - f"single value {lo} — the declared range [{d.min}, {d.max}] " - f"is advertised but the artifact admits ONE shape") + f"{input_name}[{d.axis}] symbol {sym} solved to " + f"[{lo * factor}, {hi * factor}] which does not cover " + f"the declared [{d.min}, {d.max}] — the artifact " + f"admits less traffic than it advertises") covered = True break - # The symbol may carry a multiple-of factor (8*s95), so compare the - # DECLARED bounds against the symbol's own solved bounds scaled by - # the factor the declaration states. - factor = max(1, int(d.multiple_of or 1)) - want_lo, want_hi = d.min // factor, d.max // factor - if lo > want_lo or hi < want_hi: + if not covered and not syms: gaps.append( - f"{d.input_name}[{d.axis}] symbol {sym} solved to " - f"[{lo * factor}, {hi * factor}] which does not cover the " - f"declared [{d.min}, {d.max}] — the artifact admits less " - f"traffic than it advertises") - covered = True - break - if not covered and not syms: - gaps.append( - f"{d.input_name}[{d.axis}] is symbolic ({text}) but carries no " - f"resolvable symbol; its admissible range is unprovable") + f"{input_name}[{d.axis}] is symbolic ({text}) but carries " + f"no resolvable symbol; its admissible range is unprovable") gaps.extend(_pinning_guards(program, declared_symbols)) return gaps @@ -1060,10 +1100,12 @@ def _export_entry( input_names = _input_names(module, args, kwargs) flat_names = flat_input_names(module, args, kwargs) + # ONE arity map for this arm: the spec builder mirrors the container + # structure with it and the gates below resolve declared names against + # the exported ones with it (pgw#993). + arities = _decl.container_arities(decl, espec, module) dynamic = dynamic_shapes_spec( - espec.dynamic, input_names, - _decl.container_arities(decl, espec, module), - ) if espec.dynamic else None + espec.dynamic, input_names, arities) if espec.dynamic else None def _full_export() -> Any: try: @@ -1089,11 +1131,11 @@ def _full_export() -> Any: if prop_probe is not None: timings["prop_probe_s"] = prop_probe - gaps = declared_range_gaps(program, espec.dynamic) + gaps = declared_range_gaps(program, espec.dynamic, arities) if gaps: raise MintRefused( f"entry {entry!r}: declared-range gate: " + "; ".join(gaps)) - lifted_gaps = lifted_input_gaps(program, espec) + lifted_gaps = lifted_input_gaps(program, espec, arities) if lifted_gaps: raise MintRefused( f"entry {entry!r}: lifted-input gate: " + "; ".join(lifted_gaps)) @@ -3026,7 +3068,10 @@ def _specialization_facts(spec: ExportSpec) -> Dict[str, Any]: return facts -def lifted_input_gaps(program: Any, spec: ExportSpec) -> List[str]: +def lifted_input_gaps( + program: Any, spec: ExportSpec, + containers: Optional[Mapping[str, int]] = None, +) -> List[str]: """Named reasons the declared lifted inputs are not actually graph inputs. #725 option 2's guarantee is structural: the adapter cannot be baked @@ -3035,19 +3080,24 @@ def lifted_input_gaps(program: Any, spec: ExportSpec) -> List[str]: means the branch was constant-folded, the same bug in a different hat" case. So the presence of every declared lifted input is proven here, on the program, before a single second of AOTI compile is spent. + + Names resolve through :func:`exported_input_names` (pgw#993), so a lifted + input declared as a ``repeat=`` container is looked up under the flattened + names export actually emits rather than refusing a sound program. """ if not spec.lifted_inputs: return [] signature = getattr(program, "graph_signature", None) user_inputs = {str(n) for n in getattr(signature, "user_inputs", ()) or ()} gaps: List[str] = [] - for name in spec.lifted_inputs: - if str(name) not in user_inputs: - gaps.append( - f"declared lifted input {name!r} is not a user input of the " - f"exported program (inputs: {sorted(user_inputs)!r}) — the " - f"adapter would not be swappable (#725 option 2)" - ) + for declared in spec.lifted_inputs: + for name in exported_input_names(str(declared), containers): + if name not in user_inputs: + gaps.append( + f"declared lifted input {name!r} is not a user input of " + f"the exported program (inputs: {sorted(user_inputs)!r}) " + f"— the adapter would not be swappable (#725 option 2)" + ) return gaps @@ -3350,6 +3400,7 @@ def credential() -> str: "emit_phase_events", "entry_graph_block", "export_program", + "exported_input_names", "shared_identity_blocks", "LIFTED_LORA_TORCH_FLOOR", "lifted_input_gaps", diff --git a/tests/test_dim_flattening_pgw993.py b/tests/test_dim_flattening_pgw993.py new file mode 100644 index 00000000..5020f77b --- /dev/null +++ b/tests/test_dim_flattening_pgw993.py @@ -0,0 +1,335 @@ +"""pgw#993 — a ``Dim`` carried by a ``repeat=`` container is RESOLVABLE. + +THE FIELD RECORD. gen-worker 0.93.2, pod `7evvazd2xplzml` (A100-SXM4-80GB), +\\$0.4655: the z-image AOT mint passed every earlier gate, entered the AOT +recipe, and then refused in `trace_graph` — exit=2, `deterministic`, four +identical attempts:: + + aot mint refused: entry 'transformer/adapter=true,cfg=true': + declared-range gate: declared dynamic dim names input 'x', which is not a + user input of the exported program (inputs: ['cap_feats_0', + 'cap_feats_1', 'lora_a', 'lora_b', 't', 'x_0', 'x_1']) + +THE MECHANISM — two SDK features that each work and could not compose. +`Input.repeat` containers (pgw#853) are FLATTENED by `torch.export` into one +positional user input per element, suffixed `_0`, `_1`, …; `Dim.carried_by` +names its input by the DECLARED name. The declared-range gate resolved the +declared name against the exported program and refused on the miss, which +made a `Dim` carried by a repeated container unsatisfiable by construction — +no declaration edit inside the vocabulary could fix it. z-image is the only +family in the fleet using the list vocabulary, so a rented GPU was the only +thing that could find it. + +THE INVARIANT. `carried_by` resolves through the SAME flattening +`dynamic_shapes_spec` mirrors the structure with: +`aot_mint.exported_input_names`, fed the arity map the example feed was built +from. One expansion rule, both consumers. Two independent spellings of one +name mapping is the defect class, so the sweep covers `lifted_input_gaps` +too. + +NO GPU IS NEEDED to hold this. The first half of this file drives the gate +against a program DOUBLE whose user inputs are the pod's, verbatim; the +second half exports for real on CPU and re-proves the whole path from a +declaration, for both arms of the fork (N=1 and N=2). +""" + +from __future__ import annotations + +from typing import Any, Dict, List, Sequence, Tuple + +import pytest + +from gen_worker import Compile, Dim, GraphClass, Input +from gen_worker.aot_contract import DynamicDim, ExportSpec +from gen_worker.aot_declaration import ( + cell_plans, container_arities, declared_inputs, +) +from gen_worker.aot_mint import ( + declared_range_gaps, dynamic_shapes_spec, exported_input_names, + lifted_input_gaps, +) + +# Section 3 exports for real; sections 1 and 2 drive doubles. Module-level, as +# the rest of this suite does it — torch is a dev dependency and CI installs it, +# so this never fires there. +torch = pytest.importorskip("torch") + +# The pod's sentence, verbatim. Not paraphrased: this file is the RED proof +# and the RED has to be the thing that was measured. +FIELD_REFUSAL = ( + "declared dynamic dim names input 'x', which is not a user input of the " + "exported program (inputs: ['cap_feats_0', 'cap_feats_1', 'lora_a', " + "'lora_b', 't', 'x_0', 'x_1'])") + + +# --------------------------------------------------------------------------- +# An ExportedProgram double: user inputs are all that this gate reads. +# --------------------------------------------------------------------------- + + +class _Expr: + """A sympy-shaped expression carrying one free symbol.""" + + def __init__(self, name: str) -> None: + self.name = name + self.free_symbols: Tuple[Any, ...] = (self,) + + def __str__(self) -> str: + return self.name + + +class _Node: + def __init__(self, expr: _Expr) -> None: + self.expr = expr + self.shape_env = None + + +class _SymInt: + def __init__(self, expr: _Expr) -> None: + self.node = _Node(expr) + + def __str__(self) -> str: + return str(self.node.expr) + + +class _Interval: + def __init__(self, lower: int, upper: int) -> None: + self.lower = lower + self.upper = upper + + +class _Val: + def __init__(self, shape: Tuple[Any, ...]) -> None: + self.shape = shape + + +class _Placeholder: + def __init__(self, name: str, val: _Val) -> None: + self.op = "placeholder" + self.name = name + self.meta: Dict[str, Any] = {"val": val} + + +class _Program: + """Enough ExportedProgram for `_placeholder_shapes` + the range checks.""" + + def __init__(self, shapes: Dict[str, Tuple[Any, ...]], + ranges: Dict[Any, _Interval]) -> None: + self.graph_signature = type( + "_Sig", (), {"user_inputs": tuple(shapes)})() + graph = type("_Graph", (), {"nodes": [ + _Placeholder(name, _Val(shape)) for name, shape in shapes.items() + ]})() + self.graph_module = type("_GM", (), {"graph": graph})() + self.range_constraints = ranges + + +def _zimage_arm(arity: int) -> _Program: + """The pod's program, parameterised by the CFG fork's resolved arity. + + N=2 reproduces `['cap_feats_0', 'cap_feats_1', 'lora_a', 'lora_b', 't', + 'x_0', 'x_1']` exactly; N=1 is the `cfg=false` arm, where torch still + suffixes — a one-element container exports as `x_0`, never `x` + (measured, and re-measured on a real export at the bottom of this file). + """ + height, width, caption = _Expr("s0"), _Expr("s1"), _Expr("s2") + shapes: Dict[str, Tuple[Any, ...]] = {} + for index in range(arity): + shapes[f"x_{index}"] = ( + 4, 1, _SymInt(height), _SymInt(width)) + shapes[f"cap_feats_{index}"] = (_SymInt(caption), 1024) + shapes["lora_a"] = (128, 64) + shapes["lora_b"] = (64, 128) + shapes["t"] = (arity,) + return _Program(shapes, { + height: _Interval(16, 128), + width: _Interval(16, 128), + caption: _Interval(1, 512), + }) + + +ZIMAGE_DIMS: Tuple[DynamicDim, ...] = ( + DynamicDim("x", 2, 16, 128, multiple_of=2, dim="H_lat"), + DynamicDim("x", 3, 16, 128, multiple_of=2, dim="W_lat"), + DynamicDim("cap_feats", 0, 1, 512, dim="T_cap"), +) + + +def _arities(arity: int) -> Dict[str, int]: + return {"x": arity, "cap_feats": arity} + + +# --------------------------------------------------------------------------- +# 1. THE REFUSAL, AND THE FIX +# --------------------------------------------------------------------------- + + +def test_the_double_reproduces_the_pods_refusal_verbatim() -> None: + """The double is faithful: without the arity map the gate still says the + sentence the mint died on, character for character.""" + gaps = declared_range_gaps(_zimage_arm(2), ZIMAGE_DIMS) + + assert FIELD_REFUSAL in gaps, gaps + + +def test_a_dim_carried_by_a_container_resolves_on_the_cfg_doubled_arm() -> None: + """RED before pgw#993: the gate had no arity map at all, so `x` was looked + up under its declared name and every z-image mint refused.""" + assert declared_range_gaps(_zimage_arm(2), ZIMAGE_DIMS, _arities(2)) == [] + + +def test_the_same_dim_resolves_on_the_single_element_arm() -> None: + """Both arms of the fork, because a fix for N=2 that breaks N=1 is not a + fix — `cfg=false` is the arm the declaration edit would have had to + sacrifice.""" + assert declared_range_gaps(_zimage_arm(1), ZIMAGE_DIMS, _arities(1)) == [] + + +def test_a_container_element_that_specialized_is_still_refused() -> None: + """The gate must check EVERY element, not merely find one. A per-element + pin is exactly the pgw#704 B2 defect wearing a container.""" + program = _zimage_arm(2) + shape = list(program.graph_module.graph.nodes[2].meta["val"].shape) + assert program.graph_module.graph.nodes[2].name == "x_1" + shape[2] = 64 + program.graph_module.graph.nodes[2].meta["val"].shape = tuple(shape) + + gaps = declared_range_gaps(program, ZIMAGE_DIMS, _arities(2)) + + assert len(gaps) == 1, gaps + assert gaps[0].startswith("x_1[2] exported as the STATIC value 64"), gaps + + +def test_a_name_that_is_genuinely_absent_still_refuses() -> None: + """The gate keeps its teeth: expansion resolves declared names, it does + not stop resolving them.""" + gaps = declared_range_gaps( + _zimage_arm(2), (DynamicDim("nope", 0, 2, 8),), _arities(2)) + + assert gaps and "not a user input" in gaps[0], gaps + gaps = declared_range_gaps( + _zimage_arm(2), (DynamicDim("x", 2, 16, 128),), {"x": 3}) + assert any("flattened element 'x_2'" in gap for gap in gaps), gaps + + +def test_a_plain_input_is_untouched_by_the_expansion() -> None: + """pgw#846: every declaration written before containers existed resolves + exactly as it did — a non-container name is its own exported name.""" + program = _Program( + {"sample": (2, 4, _SymInt(_Expr("s9")), 64)}, + {}) + program.range_constraints = { + program.graph_module.graph.nodes[0].meta["val"].shape[2].node.expr: + _Interval(8, 32)} + dims = (DynamicDim("sample", 2, 8, 32),) + + assert declared_range_gaps(program, dims) == [] + assert declared_range_gaps(program, dims, {"x": 2}) == [] + + +# --------------------------------------------------------------------------- +# 2. ONE EXPANSION RULE, SHARED BY BOTH CONSUMERS +# --------------------------------------------------------------------------- + + +def test_the_expansion_rule_is_the_one_torch_uses() -> None: + assert exported_input_names("x") == ("x",) + assert exported_input_names("x", {}) == ("x",) + assert exported_input_names("x", {"other": 2}) == ("x",) + assert exported_input_names("x", {"x": 1}) == ("x_0",) + assert exported_input_names("x", {"x": 3}) == ("x_0", "x_1", "x_2") + + +def test_the_spec_builder_and_the_gate_expand_identically() -> None: + """The invariant, asserted rather than described: the structure + `dynamic_shapes_spec` mirrors and the names the gate resolves come out of + the SAME call. Two spellings of one mapping is what pgw#993 was.""" + dims = (DynamicDim("x", 2, 8, 16, dim="H_lat"),) + for arity in (1, 2, 5): + spec = dynamic_shapes_spec(dims, ["x", "t"], {"x": arity}) + assert isinstance(spec["x"], list) + assert len(spec["x"]) == len(exported_input_names("x", {"x": arity})) + + +def test_a_lifted_input_declared_as_a_container_resolves_too() -> None: + """The sweep (pgw#993 acceptance 4): the lifted-input gate resolved names + the same wrong way. Nothing declares a repeated adapter today, which is + precisely the argument — z-image's `carried_by` had no user either, until + it cost \\$0.4655 to find out.""" + spec = ExportSpec( + family="harness", target="transformer", weight_lane="", + precision="bf16", lifted_inputs=("lora_a",)) + program = _Program({"lora_a_0": (128, 64), "lora_a_1": (128, 64)}, {}) + + assert lifted_input_gaps(program, spec, {"lora_a": 2}) == [] + gaps = lifted_input_gaps(program, spec) + assert gaps and "not a user input" in gaps[0], gaps + + +# --------------------------------------------------------------------------- +# 3. THE SAME THING ON A REAL EXPORT (CPU, no GPU) +# --------------------------------------------------------------------------- + + +class _ListModule(torch.nn.Module): + """z-image's shape: a python LIST of per-sample tensors.""" + + def __init__(self) -> None: + super().__init__() + self.config = type("_Cfg", (), {"in_channels": 4})() + + def forward(self, x: List[Any], t: Any) -> Any: + return torch.stack([e.sum() for e in x]) + t.sum() + + +def _list_declaration(repeat: Any) -> Compile: + return Compile( + family="harness-pgw993", + targets=("transformer",), + text_len=0, + shapes=((64, 64),), + dims=( + Dim("H_lat", carried_by=(("x", 2),)), + Dim("W_lat", carried_by=(("x", 3),)), + ), + classes=( + GraphClass(dims={"H_lat": 8, "W_lat": 8}), + GraphClass(dims={"H_lat": 16, "W_lat": 12}), + ), + inputs=( + Input("x", shape=(("config", "in_channels"), 1, "H_lat", "W_lat"), + repeat=repeat), + Input("t", shape=(2,), dtype="float32"), + ), + shape_strategy="dynamic-collapse", + warm_changes_key=False, + ) + + +def _export(arity: int) -> Tuple[Any, Sequence[DynamicDim], Dict[str, int]]: + decl = _list_declaration(arity) + spec = ExportSpec(family=decl.family, target="transformer", + weight_lane="", precision="bf16") + module = _ListModule() + args, kwargs = declared_inputs(module, spec, decl) + (plan,) = cell_plans(decl) + arities = container_arities(decl, spec, module) + program = torch.export.export( + module, tuple(args), dict(kwargs), strict=True, + dynamic_shapes=dynamic_shapes_spec(plan.dynamic, ["x", "t"], arities)) + return program, plan.dynamic, arities + + +@pytest.mark.parametrize("arity", [1, 2]) +def test_a_real_export_flattens_and_the_gate_follows(arity: int) -> None: + """End to end from a declaration, on both arms: torch really does emit + `x_0`/`x_1`, and the mint's own gate passes the program it produced.""" + program, dims, arities = _export(arity) + + assert list(program.graph_signature.user_inputs) == \ + [f"x_{i}" for i in range(arity)] + ["t"] + assert arities == {"x": arity} + assert declared_range_gaps(program, dims, arities) == [] + # And the declared-name spelling is the refusal the pod measured. + assert any("not a user input" in gap + for gap in declared_range_gaps(program, dims))