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..9df097f 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. + """ 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..7a7f348 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,218 @@ 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, +# while otherwise passing semantic validation. Without a guard the conversion +# would escape as a bare `OverflowError` instead of a categorized failure. +# -------------------------------------------------------------------------- + + +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. 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: + """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.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)), + ) + ) + + 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_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 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)