From f202d41c72e21a1c776e3354e74af4ff2ab4e840 Mon Sep 17 00:00:00 2001 From: AlekseiChirkov Date: Fri, 28 Aug 2026 13:04:33 +0700 Subject: [PATCH 1/3] test(autodiff): cover four guards the combined passes are missing Each case was reproduced against the current code before being written, so none is speculative. The constant reader returns a descriptor that can disagree with the node's own declared type, though the declaration is documented as the single authority a consumer materializes from. The agreement test could not catch it: it built both sides from one dict. It now spells them independently and asserts a disagreeing declaration is rejected. A mean whose operand declares dimensions too large to convert passes every clause of the supported-reduction predicate and then raises a bare builtin from the reciprocal. A large-but-convertible count is pinned as still expanding, so the boundary is explicit rather than implied. The two passes disagree about which identifiers a minted one must avoid. The gradient pass indexes values a program merely reads; the forward pass indexes only produced values, so a minted identifier can alias a value a node reads and silently rewire that consumer instead of failing closed. Four cases pin both passes against both positions. Two further cases pin the handler requirement the documentation calls the most likely omission: a rank-reducing derivative program lowered without a reshape handler must fail closed on the seed reshape the rewrite leaves in place. That behaviour is already correct and was covered by nothing. --- ...autodiff_expansion_lowering_equivalence.py | 44 ++++ py/tests/test_autodiff_fill_contract.py | 106 ++++++++- py/tests/test_autodiff_mean_expansion.py | 214 +++++++++++++++++- 3 files changed, 361 insertions(+), 3 deletions(-) diff --git a/py/tests/test_autodiff_expansion_lowering_equivalence.py b/py/tests/test_autodiff_expansion_lowering_equivalence.py index 52fd07a..6dc3c7f 100644 --- a/py/tests/test_autodiff_expansion_lowering_equivalence.py +++ b/py/tests/test_autodiff_expansion_lowering_equivalence.py @@ -259,6 +259,50 @@ def test_rank_reducing_expanded_graph_without_a_reshape_handler_fails_closed_nam assert unaware.invocations == [] +def test_rank_reducing_expanded_derivative_program_without_a_reshape_handler_fails_closed(): + """The documented most-likely omission, on the artifact it is easiest to miss. + + Section 9.3 says the rank-reducing tier needs a trivial-reshape handler in + *both* artifacts. On the gradient path the reshape is not emitted by the + expansion at all -- it is the seed reshape the VJP already produced, which + the broadcast-and-scale rewrite leaves in place -- so a backend that added a + fill handler and stopped is rejected by that surviving node. + """ + graph, program = _traced_mean(keepdims=False) + expanded = expand_mean_derivative_program(program) + reshape_node_ids = [ + node.node_id for node in expanded.nodes if isinstance(node.operator, ReshapeOperator) + ] + assert reshape_node_ids, "the rank-reducing gradient path must retain a reshape node" + unaware = recording_registry(_limited(reshape=False)) + assert unaware.registry.has_handler(FillOperator()) + assert unaware.registry.has_handler(MatmulOperator()) + values = {"seed": _seed(keepdims=False, dtype="f64"), "v0": _operand((3, 5), "f64")} + + with pytest.raises(AutodiffError) as excinfo: + _lower_gradient( + expanded, forward_graph=graph, registry=unaware.registry, values=values + ) + + assert excinfo.value.category == "unsupported_operator" + assert "ReshapeOperator" in excinfo.value.message + assert reshape_node_ids[0] in excinfo.value.message + assert unaware.invocations == [] + + +def test_rank_reducing_expanded_derivative_program_lowers_once_a_reshape_handler_exists(): + """The positive control: the same artifact and registry, reshape handler added.""" + graph, program = _traced_mean(keepdims=False) + expanded = expand_mean_derivative_program(program) + values = {"seed": _seed(keepdims=False, dtype="f64"), "v0": _operand((3, 5), "f64")} + + gradient = _lower_gradient( + expanded, forward_graph=graph, registry=_limited(reshape=True), values=values + ) + + np.testing.assert_allclose(gradient, np.full((3, 5), 2.5 / 15.0), rtol=1e-12) + + def test_registry_lookup_message_names_only_the_operator_type(): """`lookup`'s own message stays byte-identical; only the pre-flight enriches it.""" with pytest.raises(AutodiffError) as excinfo: diff --git a/py/tests/test_autodiff_fill_contract.py b/py/tests/test_autodiff_fill_contract.py index 4fd3c96..3d5973e 100644 --- a/py/tests/test_autodiff_fill_contract.py +++ b/py/tests/test_autodiff_fill_contract.py @@ -143,15 +143,119 @@ def test_fill_descriptor_accepts_a_zero_dimensional_shape() -> None: def test_fill_node_output_typespec_equals_its_descriptor_dtype_and_shape() -> None: + """A fill node's declaration and its parameters agree, and the reader enforces it. + + The two sides are written out independently here rather than derived from + one dict, and the negative half asserts that a declaration disagreeing with + the parameters is rejected -- so this test fails if the reader stops + cross-checking the declaration it is the single authority for (risk R-3). + """ from tinychain.autodiff import fill_descriptor - node = _fill_node() + node = _fill_node( + op_params={"fill": 0.25, "dtype": "f32", "shape": [2, 3]}, + output_typespec={"dtype": "f32", "shape": [2, 3]}, + ) descriptor = fill_descriptor(node) assert node.output_typespec is not None assert node.output_typespec["dtype"] == descriptor.dtype assert tuple(node.output_typespec["shape"]) == descriptor.shape + disagreeing = _fill_node( + op_params={"fill": 0.25, "dtype": "f32", "shape": [2, 3]}, + output_typespec={"dtype": "f64", "shape": [9, 9]}, + ) + with pytest.raises(AutodiffError): + fill_descriptor(disagreeing) + + +def test_fill_descriptor_rejects_a_node_declaring_a_disagreeing_dtype() -> None: + from tinychain.autodiff import fill_descriptor + + node = _fill_node( + node_id="fill_dtype_disagrees", + op_params={"fill": 1.0, "dtype": "f64", "shape": [2, 2]}, + output_typespec={"dtype": "f32", "shape": [2, 2]}, + ) + + with pytest.raises(AutodiffError) as raised: + fill_descriptor(node) + + assert raised.value.category == "malformed_derivative_ir" + assert "fill_dtype_disagrees" in raised.value.message + assert "f32" in raised.value.message + assert "f64" in raised.value.message + + +def test_fill_descriptor_rejects_a_node_declaring_a_disagreeing_shape() -> None: + from tinychain.autodiff import fill_descriptor + + node = _fill_node( + node_id="fill_shape_disagrees", + op_params={"fill": 1.0, "dtype": "f64", "shape": [2, 2]}, + output_typespec={"dtype": "f64", "shape": [9, 9]}, + ) + + with pytest.raises(AutodiffError) as raised: + fill_descriptor(node) + + assert raised.value.category == "malformed_derivative_ir" + assert "fill_shape_disagrees" in raised.value.message + assert "[9, 9]" in raised.value.message or "(9, 9)" in raised.value.message + + +def test_fill_descriptor_rejects_an_operation_context_declaring_a_disagreeing_type() -> None: + """The reader is the same authority over the context a handler receives.""" + from tinychain.autodiff import fill_descriptor + + node = _fill_node( + node_id="fill_context_disagrees", + op_params={"fill": 1.0, "dtype": "f64", "shape": [2, 2]}, + output_typespec={"dtype": "f32", "shape": [9, 9]}, + ) + + with pytest.raises(AutodiffError) as raised: + fill_descriptor(_operation_context(node)) + + assert raised.value.category == "malformed_derivative_ir" + assert "fill_context_disagrees" in raised.value.message + + +def test_fill_descriptor_reads_a_fill_node_that_declares_no_typespec() -> None: + """An absent declaration stays acceptable: the reader validates parameters, + and cross-checks a declaration only when one is present.""" + from tinychain.autodiff import fill_descriptor + + from tinychain.autodiff import FillOperator + + node = TensorNodeRecord( + node_id="fill_undeclared", + output_value_id="fill_undeclared_out", + operator=FillOperator(), + op_params=_well_formed_params(), + input_value_ids=[], + ) + assert node.output_typespec is None + + descriptor = fill_descriptor(node) + + assert descriptor.dtype == "f32" + assert descriptor.shape == (2, 3) + assert fill_descriptor(_operation_context(node)) == descriptor + + +def test_fill_descriptor_accepts_a_declaration_spelled_with_a_tuple_shape() -> None: + """Agreement is about the dimensions, not the container the shape uses.""" + from tinychain.autodiff import fill_descriptor + + node = _fill_node( + op_params={"fill": 1.0, "dtype": "f64", "shape": [2, 3]}, + output_typespec={"dtype": "f64", "shape": (2, 3)}, + ) + + assert fill_descriptor(node).shape == (2, 3) + # -------------------------------------------------------------------------- # Categorized descriptor failures diff --git a/py/tests/test_autodiff_mean_expansion.py b/py/tests/test_autodiff_mean_expansion.py index 06d542b..6f4c0c6 100644 --- a/py/tests/test_autodiff_mean_expansion.py +++ b/py/tests/test_autodiff_mean_expansion.py @@ -16,8 +16,12 @@ * each supported-mean validation failure raises its own category with a message naming the offending node and explaining the failed condition. -Nothing here asserts anything about the gradient-path rewrite, provenance -records, or the detailed passes; those are separate work. +Nothing here asserts anything about the gradient-path *rewrite*, provenance +records, or the detailed passes; those are separate work. The one exception is +the identifier-collision section at the end: one helper indexes the value ids an +artifact mentions for both passes, so the cases proving a minted identifier can +never alias an existing value are written against both passes together, where +the shared contract can be seen as one thing. """ from __future__ import annotations @@ -790,3 +794,209 @@ def test_expand_mean_graph_is_exported_from_the_autodiff_package() -> None: assert "expand_mean_graph" in autodiff.__all__ assert callable(autodiff.expand_mean_graph) assert not hasattr(tc, "expand_mean_graph") + + +# -------------------------------------------------------------------------- +# an operand whose declared element count cannot be represented +# +# The reciprocal `1 / (rows * columns)` is the one inexact substitution the +# region makes. A declared shape can name an element count no float can hold, +# and such a mean passes every clause of the predicate -- so without a guard +# the conversion escapes as a bare `OverflowError`, which NFR-128-004 forbids. +# -------------------------------------------------------------------------- + + +def test_a_mean_whose_element_count_cannot_be_converted_is_rejected() -> None: + rows = columns = 10**200 + graph = _mean_graph(node_id="huge", operand_shape=(rows, columns)) + + with pytest.raises(AutodiffError) as raised: + _expand(graph) + + assert raised.value.category == "unsupported_reduction" + assert "huge" in raised.value.message + assert str(rows * columns) in raised.value.message + + +def test_a_mean_whose_element_count_cannot_be_converted_raises_no_bare_builtin() -> None: + graph = _mean_graph(node_id="huge", operand_shape=(10**200, 10**200)) + + try: + _expand(graph) + except AutodiffError: + pass + except (OverflowError, KeyError, IndexError, TypeError, ValueError) as exc: + pytest.fail(f"bare {type(exc).__name__} escaped expand_mean_graph: {exc}") + + +def test_a_large_but_convertible_element_count_still_expands() -> None: + """The guard rejects only what cannot be converted; a merely huge count expands.""" + rows = columns = 10**150 + graph = _mean_graph(operand_shape=(rows, columns)) + + expanded = _expand(graph) + + scale = expanded.nodes[4] + assert isinstance(scale.operator, MulOperator) + assert scale.op_params["right_literal"] == 1.0 / float(rows * columns) + + +# -------------------------------------------------------------------------- +# a minted identifier can never alias a value the artifact merely mentions +# +# Both passes mint into the reserved namespace and both must fail closed on a +# collision (Inv-5). Indexing only *produced* values is not enough: a minted id +# equal to one a node merely reads, or one named only among the artifact's +# declared outputs, would silently rewire that consumer to the emitted node -- +# a wrong result with no error at all. +# -------------------------------------------------------------------------- + + +def _graph_with_a_mean_and( + *, extra_nodes: list[TensorNodeRecord] | None = None, extra_outputs: list[str] | None = None +) -> TensorGraph: + """A traced supported mean, plus whatever extra nodes or outputs a case needs.""" + graph = _traced_mean_graph(keepdims=True) + return TensorGraph( + nodes=[*graph.nodes, *(extra_nodes or [])], + inputs=list(graph.inputs), + outputs=[*graph.outputs, *(extra_outputs or [])], + ) + + +def _reading_node(value_id: str) -> TensorNodeRecord: + """A node that reads *value_id* and nothing else produces it.""" + return TensorNodeRecord( + node_id="n9", + output_value_id="v9", + operator=SumOperator(), + op_params={"axes": [0], "keepdims": True}, + input_value_ids=[value_id], + output_typespec=_typespec("f64", (1, 1)), + ) + + +def _broadcast_scale_program(*, extra_output_gradient: str | None = None): + """One rewritable broadcast-and-scale chain, minimal and hand-built. + + Built here rather than imported so this module keeps its own inputs; the + chain is the smallest one satisfying the gradient-path predicate, because + the pass must match a region before it mints anything at all. + """ + from tinychain.autodiff import ( + BroadcastOperator, + DerivativeMetadata, + DerivativeProgram, + DivOperator, + ) + + rows, columns = 3, 5 + source = _typespec("f64", (1, 1)) + broadcast_typespec = _typespec("f64", (rows, columns)) + broadcast = TensorNodeRecord( + node_id="dn0", + output_value_id="d0", + operator=BroadcastOperator(), + op_params={"shape": [rows, columns]}, + input_value_ids=["seed"], + output_typespec=broadcast_typespec, + ) + division = TensorNodeRecord( + node_id="dn1", + output_value_id="d1", + operator=DivOperator(), + op_params={"right_literal": float(rows * columns)}, + input_value_ids=["d0"], + output_typespec=broadcast_typespec, + ) + output_gradients: list[str | None] = ["d1"] + if extra_output_gradient is not None: + output_gradients.append(extra_output_gradient) + return DerivativeProgram( + nodes=[broadcast, division], + gradients={"v0": "d1"}, + output_gradients=output_gradients, + metadata=DerivativeMetadata( + source_graph_id="graph", + transform_version="0.1.0", + tensor_op_contract_version="0.1.0", + wrt_signature=("v0",), + seed_contract="seed matches output", + ), + value_typespecs={"seed": source, "d0": broadcast_typespec, "d1": broadcast_typespec}, + ) + + +def _expand_program(program): + from tinychain.autodiff import expand_mean_derivative_program + + return expand_mean_derivative_program(program) + + +def test_the_forward_pass_rejects_a_minted_value_id_a_node_merely_reads() -> None: + """The consumer must not be silently rewired to the emitted constant.""" + reserved = _reserved_value_id() + seeded = _graph_with_a_mean_and( + extra_nodes=[_reading_node(reserved)], extra_outputs=["v9"] + ) + + with pytest.raises(AutodiffError) as raised: + _expand(seeded) + + assert raised.value.category == "malformed_derivative_ir" + assert reserved in raised.value.message + + +def test_the_forward_pass_rejects_a_minted_value_id_named_only_in_the_outputs() -> None: + reserved = _reserved_value_id() + seeded = _graph_with_a_mean_and(extra_outputs=[reserved]) + + with pytest.raises(AutodiffError) as raised: + _expand(seeded) + + assert raised.value.category == "malformed_derivative_ir" + assert reserved in raised.value.message + + +def test_the_forward_pass_never_rewires_a_consumer_to_an_emitted_node() -> None: + """The fail-closed half stated as the property it protects: whatever value a + pre-existing node read before the pass, it still reads afterwards.""" + graph = _traced_mean_graph(keepdims=True) + reads_before = {node.node_id: list(node.input_value_ids) for node in graph.nodes} + + expanded = _expand(graph) + + for node in expanded.nodes: + if node.node_id in reads_before: + assert list(node.input_value_ids) == reads_before[node.node_id] + + +def test_the_gradient_pass_rejects_a_minted_value_id_a_node_merely_reads() -> None: + reserved = _reserved_value_id() + program = _broadcast_scale_program() + program.nodes[0] = TensorNodeRecord( + node_id="dn0", + output_value_id="d0", + operator=program.nodes[0].operator, + op_params=dict(program.nodes[0].op_params), + input_value_ids=[reserved], + output_typespec=program.nodes[0].output_typespec, + ) + program.value_typespecs.pop("seed") + + with pytest.raises(AutodiffError) as raised: + _expand_program(program) + + assert raised.value.category == "malformed_derivative_ir" + assert reserved in raised.value.message + + +def test_the_gradient_pass_rejects_a_minted_value_id_named_only_in_the_output_gradients() -> None: + reserved = _reserved_value_id() + program = _broadcast_scale_program(extra_output_gradient=reserved) + + with pytest.raises(AutodiffError) as raised: + _expand_program(program) + + assert raised.value.category == "malformed_derivative_ir" + assert reserved in raised.value.message From 8896de194d8a9da028aeac9923b781e763ccf723 Mon Sep 17 00:00:00 2001 From: AlekseiChirkov Date: Fri, 28 Aug 2026 13:10:05 +0700 Subject: [PATCH 2/3] test(autodiff): give the collision case a bystander consumer The gradient read-position case fed the reserved identifier into the chain it was expanding, which stopped the region matching, so nothing was minted and the collision it asserts could never fire. Read the identifier from a node outside the region instead. Assertions unchanged. --- py/tests/test_autodiff_mean_expansion.py | 29 ++++++++++++++++-------- 1 file changed, 19 insertions(+), 10 deletions(-) diff --git a/py/tests/test_autodiff_mean_expansion.py b/py/tests/test_autodiff_mean_expansion.py index 6f4c0c6..f945a9f 100644 --- a/py/tests/test_autodiff_mean_expansion.py +++ b/py/tests/test_autodiff_mean_expansion.py @@ -853,7 +853,9 @@ def test_a_large_but_convertible_element_count_still_expands() -> None: def _graph_with_a_mean_and( - *, extra_nodes: list[TensorNodeRecord] | None = None, extra_outputs: list[str] | None = None + *, + extra_nodes: list[TensorNodeRecord] | None = None, + extra_outputs: list[str] | None = None, ) -> TensorGraph: """A traced supported mean, plus whatever extra nodes or outputs a case needs.""" graph = _traced_mean_graph(keepdims=True) @@ -972,17 +974,24 @@ def test_the_forward_pass_never_rewires_a_consumer_to_an_emitted_node() -> None: def test_the_gradient_pass_rejects_a_minted_value_id_a_node_merely_reads() -> None: + """The reserved id is read by a bystander node, so the chain still matches. + + Feeding it to the chain itself would only prove the predicate declines an + unmatched region -- the pass would mint nothing at all, and the collision + this case is about would never arise. + """ reserved = _reserved_value_id() program = _broadcast_scale_program() - program.nodes[0] = TensorNodeRecord( - node_id="dn0", - output_value_id="d0", - operator=program.nodes[0].operator, - op_params=dict(program.nodes[0].op_params), - input_value_ids=[reserved], - output_typespec=program.nodes[0].output_typespec, + program.nodes.append( + TensorNodeRecord( + node_id="dn9", + output_value_id="d9", + operator=MulOperator(), + op_params={"right_literal": 2.0}, + input_value_ids=[reserved], + output_typespec=_typespec("f64", (3, 5)), + ) ) - program.value_typespecs.pop("seed") with pytest.raises(AutodiffError) as raised: _expand_program(program) @@ -991,7 +1000,7 @@ def test_the_gradient_pass_rejects_a_minted_value_id_a_node_merely_reads() -> No assert reserved in raised.value.message -def test_the_gradient_pass_rejects_a_minted_value_id_named_only_in_the_output_gradients() -> None: +def test_the_gradient_pass_rejects_a_minted_id_named_only_in_the_output_gradients() -> None: reserved = _reserved_value_id() program = _broadcast_scale_program(extra_output_gradient=reserved) From 18292eec7584298e4ea52da1f85de284d8da6cef Mon Sep 17 00:00:00 2001 From: AlekseiChirkov Date: Fri, 28 Aug 2026 13:10:05 +0700 Subject: [PATCH 3/3] fix(autodiff): close four gaps between the two expansion passes The constant reader now enforces the agreement its node's declaration is documented to be the authority for. A descriptor disagreeing with the declared dtype or shape is rejected naming the node, so a consumer cannot materialize one tensor while the graph declares another. The check runs after parameter validation, so every existing rejection keeps its category, and a node with no declaration still reads. An operand declaring dimensions too large to convert now fails validation naming the node and the count, rather than raising a bare builtin from the reciprocal. Such a shape violates no clause -- it is a positive integer -- so this is a limit on what the expansion can express, which is what the residual category covers. The threshold is the conversion limit, not the exactness limit: a count above the latter still expands to well within the documented tolerance, and refusing it would narrow the supported domain rather than guard it. Both passes now derive collision candidates from one helper covering every identifier an artifact mentions -- produced, read, declared as an input, and declared as an output. Previously each half indexed a different set, so the forward pass could mint an identifier a node already read and silently rewire that consumer to a generated constant instead of failing closed. Sharing one helper removes the contradiction rather than correcting it in two places. --- py/tests/test_autodiff_fill_contract.py | 2 +- py/tests/test_autodiff_mean_expansion.py | 6 +- py/tinychain/autodiff/expansion.py | 169 ++++++++++++++++++++--- 3 files changed, 154 insertions(+), 23 deletions(-) diff --git a/py/tests/test_autodiff_fill_contract.py b/py/tests/test_autodiff_fill_contract.py index 3d5973e..9df097f 100644 --- a/py/tests/test_autodiff_fill_contract.py +++ b/py/tests/test_autodiff_fill_contract.py @@ -148,7 +148,7 @@ def test_fill_node_output_typespec_equals_its_descriptor_dtype_and_shape() -> No The two sides are written out independently here rather than derived from one dict, and the negative half asserts that a declaration disagreeing with the parameters is rejected -- so this test fails if the reader stops - cross-checking the declaration it is the single authority for (risk R-3). + cross-checking the declaration it is the single authority for. """ from tinychain.autodiff import fill_descriptor diff --git a/py/tests/test_autodiff_mean_expansion.py b/py/tests/test_autodiff_mean_expansion.py index f945a9f..7a7f348 100644 --- a/py/tests/test_autodiff_mean_expansion.py +++ b/py/tests/test_autodiff_mean_expansion.py @@ -801,8 +801,8 @@ def test_expand_mean_graph_is_exported_from_the_autodiff_package() -> None: # # The reciprocal `1 / (rows * columns)` is the one inexact substitution the # region makes. A declared shape can name an element count no float can hold, -# and such a mean passes every clause of the predicate -- so without a guard -# the conversion escapes as a bare `OverflowError`, which NFR-128-004 forbids. +# while otherwise passing semantic validation. Without a guard the conversion +# would escape as a bare `OverflowError` instead of a categorized failure. # -------------------------------------------------------------------------- @@ -845,7 +845,7 @@ def test_a_large_but_convertible_element_count_still_expands() -> None: # a minted identifier can never alias a value the artifact merely mentions # # Both passes mint into the reserved namespace and both must fail closed on a -# collision (Inv-5). Indexing only *produced* values is not enough: a minted id +# collision. Indexing only *produced* values is not enough: a minted id # equal to one a node merely reads, or one named only among the artifact's # declared outputs, would silently rewire that consumer to the emitted node -- # a wrong result with no error at all. diff --git a/py/tinychain/autodiff/expansion.py b/py/tinychain/autodiff/expansion.py index 8795bce..3a9c9b7 100644 --- a/py/tinychain/autodiff/expansion.py +++ b/py/tinychain/autodiff/expansion.py @@ -63,6 +63,12 @@ node carries an operand, ``op_params`` holds a key outside the schema, or ``fill`` is absent or not a real number. A ``bool`` is not a real number here: ``True`` reaching a numeric field is a construction defect, not the value one. + Also raised when a node *declares* an ``output_typespec`` disagreeing with its + own parameters: the declaration is the single authority for what the node + produces, so a node promising one tensor and describing another is malformed + rather than merely inconsistent. A node carrying no declaration at all is not + a defect -- the reader validates parameters, and cross-checks a declaration + only when one is present. * ``missing_shape_metadata`` -- ``shape`` is absent, is not a sequence of integers, or holds a negative or symbolic dimension. * ``missing_dtype_metadata`` -- ``dtype`` is absent or is not a non-empty string. @@ -281,7 +287,7 @@ from __future__ import annotations -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass from .graph import ( @@ -344,6 +350,11 @@ def fill_descriptor(operation: TensorNodeRecord | OperationContext) -> FillDescr :class:`OperationContext` a lowering handler receives for the same node. Raises a categorized :class:`AutodiffError` naming the node for any node that is not a well-formed fill. + + When the node carries an ``output_typespec``, the declaration is required to + agree with the parameters. Without that check the node would declare one + tensor while its descriptor produced another, and a `Fill` handler and a + `Matmul` handler reading the same node would disagree about the constant. """ if not isinstance(operation, (TensorNodeRecord, OperationContext)): raise AutodiffError( @@ -359,11 +370,13 @@ def fill_descriptor(operation: TensorNodeRecord | OperationContext) -> FillDescr op_params = operation.op_params _check_param_keys(node_id, op_params) - return FillDescriptor( + descriptor = FillDescriptor( fill=_read_fill(node_id, op_params), dtype=_read_dtype(node_id, op_params), shape=_read_shape(node_id, op_params), ) + _check_declaration_agrees(node_id, operation.output_typespec, descriptor) + return descriptor def _check_operator(node_id: str, operator: TensorOperator) -> None: @@ -451,6 +464,55 @@ def _read_shape(node_id: str, op_params: Mapping[str, object]) -> tuple[int, ... return tuple(dimensions) +def _check_declaration_agrees( + node_id: str, output_typespec: object, descriptor: FillDescriptor +) -> None: + """Require a declared `output_typespec` to agree with the parameters read. + + A fill node's own complete `output_typespec` is the single authority for + what the node produces, so a declaration disagreeing with the parameters is + a structural defect rather than a preference: one handler would materialize + the descriptor's tensor while the graph promised the declaration's. + + An **absent** declaration is not a defect. This reader validates parameters; + it cross-checks a declaration only when the node carries one, so a node + built without a typespec still reads successfully. + """ + if output_typespec is None: + return + if not isinstance(output_typespec, Mapping): + raise AutodiffError( + "malformed_derivative_ir", + f"fill node {node_id!r} declares an output_typespec that is not a mapping: " + f"{output_typespec!r}", + ) + + declared_dtype = output_typespec.get("dtype") + if declared_dtype != descriptor.dtype: + raise AutodiffError( + "malformed_derivative_ir", + f"fill node {node_id!r} declares dtype {declared_dtype!r}, but its parameters " + f"describe {descriptor.dtype!r}; a fill node's declaration is the single " + "authority for what it produces and must agree with its own descriptor", + ) + + declared_shape = output_typespec.get("shape") + if isinstance(declared_shape, (str, bytes)) or not isinstance(declared_shape, Sequence): + raise AutodiffError( + "malformed_derivative_ir", + f"fill node {node_id!r} declares no ranked shape in its output_typespec: " + f"{declared_shape!r}", + ) + if tuple(declared_shape) != descriptor.shape: + raise AutodiffError( + "malformed_derivative_ir", + f"fill node {node_id!r} declares shape {list(declared_shape)!r}, but its " + f"parameters describe {list(descriptor.shape)!r}; a fill node's declaration " + "is the single authority for what it produces and must agree with its own " + "descriptor", + ) + + # -------------------------------------------------------------------------- # The reserved identifier namespace # -------------------------------------------------------------------------- @@ -612,6 +674,46 @@ def _indexed_value_ids( return frozenset(seen) +def _mentioned_value_ids( + nodes: Sequence[TensorNodeRecord], + *, + produced_inputs: Sequence[tuple[str, object]] = (), + declared_value_ids: Iterable[object] = (), +) -> frozenset[str]: + """Return every value id the artifact mentions, in any position. + + Wider than the set of *produced* values on purpose: a minted identifier must + not collide with a value a node merely reads, one the artifact declares as + an output, or one it merely records a typespec for either. + + Indexing only produced values is not a smaller guard, it is a silent one. A + minted id equal to one an existing node reads does not collide with anything + the index can see, so nothing is rejected -- and that node is quietly rewired + to read the emitted constant instead of the value it was written against. + The result would be a wrong artifact with no error, so the pass must reject + the collision. + + Both passes share this one index. The two halves of this module previously + held contradictory contracts here; a correct copy in a second place would + have preserved the drift rather than removed it. + + *produced_inputs* are `(value id, typespec)` boundary declarations whose ids + are produced, so they take part in duplicate detection alongside node + outputs. *declared_value_ids* are ids the artifact merely names -- declared + outputs, gradient results, inherited typespec keys -- which are indexed + against collision but never checked for duplication, because naming a value + twice in those positions is not a defect. A non-string entry (an absent + output gradient is spelled ``None``) names no value and is skipped. + """ + value_ids = set(_indexed_value_ids(nodes, produced_inputs)) + value_ids.update( + value_id for value_id in declared_value_ids if isinstance(value_id, str) + ) + for node in nodes: + value_ids.update(node.input_value_ids) + return frozenset(value_ids) + + def _declared_typespecs( nodes: Sequence[TensorNodeRecord], inputs: Sequence[tuple[str, object]] ) -> dict[str, object]: @@ -642,6 +744,7 @@ class _SupportedMean: dtype: str rows: int columns: int + reciprocal_count: float keepdims: bool output_value_id: str output_typespec: dict[str, object] @@ -711,6 +814,35 @@ def _mean_operand_dimensions(node_id: str, shape: Shape) -> tuple[int, int]: return int(rows), int(columns) +def _mean_reciprocal_count(node_id: str, rows: int, columns: int) -> float: + """Return `1 / (rows * columns)`, rejecting a count no float can represent. + + The reciprocal is the one inexact substitution the emitted region makes, and + it is the only place the pass converts a declared dimension to a float. A + shape may name an element count larger than any float while remaining + otherwise well formed because the dimension is still a positive integer. + Such a failure is a limit of what the region can express, so it is reported + as `unsupported_reduction` instead of escaping as a bare `OverflowError`. + + The boundary is representability, not exactness. A count above `2 ** 53` + converts with a relative error near 1e-16, well inside the tolerances of + 1e-6 and 1e-12, and such a mean expands correctly today; rejecting it would + narrow the supported domain rather than guard it. + """ + element_count = rows * columns + try: + count_as_float = float(element_count) + except OverflowError as exc: + raise AutodiffError( + "unsupported_reduction", + f"mean node {node_id!r} declares an operand of {rows!r} x {columns!r}, whose " + f"element count {element_count!r} is too large to convert to a float; the " + "expanded region scales by the reciprocal of that count, which this pass " + "cannot represent", + ) from exc + return 1.0 / count_as_float + + def _mean_axes(node: TensorNodeRecord) -> tuple[int, ...]: """Require the declared axes to normalize to every operand axis.""" axes = node.op_params.get("axes") @@ -830,6 +962,7 @@ def _supported_mean( operand_shape, operand_dtype = _mean_operand_shape_and_dtype(node, typespecs) rows, columns = _mean_operand_dimensions(node.node_id, operand_shape) + reciprocal_count = _mean_reciprocal_count(node.node_id, rows, columns) axes = _mean_axes(node) keepdims = _mean_keepdims(node) dtype = _mean_dtype(node.node_id, operand_dtype) @@ -845,6 +978,7 @@ def _supported_mean( dtype=dtype, rows=rows, columns=columns, + reciprocal_count=reciprocal_count, keepdims=keepdims, output_value_id=node.output_value_id, output_typespec=output_typespec, @@ -979,7 +1113,9 @@ def _emit_mean_region( row_ones = _emit_fill(minter, dtype=supported.dtype, shape=(1, rows)) total_sum = _emit_matmul(minter, left=row_ones, right=row_sums) - reciprocal_count = 1.0 / float(rows * columns) + # Proven convertible during validation, so the region can never be reached + # with a count the scale cannot express. + reciprocal_count = supported.reciprocal_count if supported.keepdims: # The rank-preserving tier: the scale already has the mean's rank-2 # shape, so it carries the mean's own value id and terminates the region. @@ -1029,7 +1165,9 @@ def expand_mean_graph_detailed(graph: TensorGraph) -> MeanGraphExpansionResult: """ nodes = graph.nodes existing_node_ids = _indexed_node_ids(nodes) - existing_value_ids = _indexed_value_ids(nodes, graph.inputs) + existing_value_ids = _mentioned_value_ids( + nodes, produced_inputs=graph.inputs, declared_value_ids=graph.outputs + ) typespecs = _declared_typespecs(nodes, graph.inputs) # Validate every candidate before emitting any replacement nodes. @@ -1307,20 +1445,6 @@ def _emit_broadcast_scale_region( # -------------------------------------------------------------------------- -def _existing_value_ids(program: DerivativeProgram) -> frozenset[str]: - """Return every value id the program names anywhere. - - Wider than the set of produced values on purpose: a minted identifier must - not collide with a value the program merely reads or merely records a - typespec for either. - """ - value_ids = set(_indexed_value_ids(program.nodes, [])) - value_ids.update(program.value_typespecs) - for node in program.nodes: - value_ids.update(node.input_value_ids) - return frozenset(value_ids) - - def expand_mean_derivative_program_detailed( program: DerivativeProgram, ) -> MeanDerivativeExpansionResult: @@ -1354,7 +1478,14 @@ def expand_mean_derivative_program_detailed( """ nodes = program.nodes existing_node_ids = _indexed_node_ids(nodes) - existing_value_ids = _existing_value_ids(program) + existing_value_ids = _mentioned_value_ids( + nodes, + declared_value_ids=( + *program.value_typespecs, + *program.gradients.values(), + *program.output_gradients, + ), + ) typespecs = _declared_typespecs(nodes, list(program.value_typespecs.items())) producers = {node.output_value_id: node for node in nodes} consumer_counts = _value_consumer_counts(nodes)