diff --git a/py/tests/test_autodiff_mean_expansion.py b/py/tests/test_autodiff_mean_expansion.py new file mode 100644 index 0000000..233c539 --- /dev/null +++ b/py/tests/test_autodiff_mean_expansion.py @@ -0,0 +1,792 @@ +"""Unit tests for the forward all-axis mean expansion pass. + +`expand_mean_graph` rewrites every supported all-axis rank-2 `MeanOperator` in a +`TensorGraph` into the matmul-based region of the specification, in both tiers: +five nodes when the mean declared `keepdims=True`, and six -- the sixth a real +`ReshapeOperator` -- when it declared `keepdims=False`. + +These tests pin three things the rewrite turns on: + +* every emitted node declares the dtype and shape its operation *actually* + produces, recomputed here from the operands rather than read back from the + node under test, so a node can never declare a shape its operation cannot + produce; +* every candidate mean is validated before any node is emitted, so a rejected + artifact never comes back partially rewritten; +* each failing clause of the supported-mean predicate raises its own category, + with a message naming both the offending node and the clause. + +Nothing here asserts anything about the gradient-path rewrite, provenance +records, or the detailed passes; those are separate work. +""" + +from __future__ import annotations + +import copy + +import pytest +import tinychain as tc +from tinychain.autodiff import ( + AutodiffError, + MatmulOperator, + MeanOperator, + MulOperator, + ReshapeOperator, + SumOperator, + TensorGraph, + TensorGraphBuilder, + TensorNodeRecord, +) + + +# -------------------------------------------------------------------------- +# lazily resolved surface under test +# +# The pass and its reserved-namespace constants are resolved inside the test +# body rather than at import time, so a missing name fails one test that has +# already built its input rather than aborting collection of the whole module. +# -------------------------------------------------------------------------- + + +def _expand(graph: TensorGraph) -> TensorGraph: + from tinychain.autodiff import expand_mean_graph + + return expand_mean_graph(graph) + + +def _reserved_node_id(index: int = 0) -> str: + from tinychain.autodiff.expansion import EXPANSION_NODE_ID_PREFIX + + return f"{EXPANSION_NODE_ID_PREFIX}{index}" + + +def _reserved_value_id(index: int = 0) -> str: + from tinychain.autodiff.expansion import EXPANSION_VALUE_ID_PREFIX + + return f"{EXPANSION_VALUE_ID_PREFIX}{index}" + + +# -------------------------------------------------------------------------- +# helpers +# -------------------------------------------------------------------------- + +_DERIVED = object() + + +def _typespec(dtype: str, shape: object) -> dict[str, object]: + return {"dtype": dtype, "shape": list(shape)} + + +def _traced_mean_graph( + *, shape: tuple[int, ...] = (3, 5), dtype: str = "f64", keepdims: bool = True +) -> TensorGraph: + """Trace `value.mean([0, 1], keepdims=...)` and return the finalized graph.""" + with TensorGraphBuilder() as trace: + value = trace.input("value", dtype=dtype, shape=shape) + output = value.mean([0, 1], keepdims=keepdims) + return trace.build(outputs=output) + + +def _mean_graph( + *, + operand_typespec: object = _DERIVED, + op_params: object = _DERIVED, + output_typespec: object = _DERIVED, + node_id: str = "n0", + output_value_id: str = "v1", + input_value_ids: list[str] | None = None, + operand_dtype: str = "f64", + operand_shape: object = (3, 5), +) -> TensorGraph: + """Build a one-node mean graph directly, with every field overridable. + + Hand construction is what lets the failure table reach malformed shapes, + dtypes, axes, and identifiers that the tracer would reject long before the + pass ever saw them. + """ + params = ( + {"axes": [0, 1], "keepdims": True} if op_params is _DERIVED else dict(op_params) # type: ignore[arg-type] + ) + resolved_operand_typespec = ( + _typespec(operand_dtype, operand_shape) + if operand_typespec is _DERIVED + else operand_typespec + ) + if output_typespec is _DERIVED: + reduced_shape: tuple[int, ...] = (1, 1) if params.get("keepdims") else () + resolved_output_typespec: object = _typespec(operand_dtype, reduced_shape) + else: + resolved_output_typespec = output_typespec + + node = TensorNodeRecord( + node_id=node_id, + output_value_id=output_value_id, + operator=MeanOperator(), + op_params=params, + input_value_ids=["v0"] if input_value_ids is None else input_value_ids, + output_typespec=resolved_output_typespec, # type: ignore[arg-type] + ) + return TensorGraph( + nodes=[node], + inputs=[("v0", resolved_operand_typespec)], # type: ignore[list-item] + outputs=[output_value_id], + ) + + +def _value_typespecs(graph: TensorGraph) -> dict[str, object]: + """Index every value id in *graph* to the typespec its producer declares.""" + typespecs: dict[str, object] = { + value_id: typespec for value_id, typespec in graph.inputs + } + for node in graph.nodes: + typespecs[node.output_value_id] = node.output_typespec + return typespecs + + +def _value_ids(graph: TensorGraph) -> set[str]: + ids = {value_id for value_id, _ in graph.inputs} + for node in graph.nodes: + ids.add(node.output_value_id) + ids.update(node.input_value_ids) + return ids + + +def _independently_computed_typespec( + node: TensorNodeRecord, typespecs: dict[str, object] +) -> dict[str, object]: + """Recompute what *node* truly produces, from its operands and its own rule. + + Deliberately written out here instead of calling the framework's shape + helpers: the point of the audit is that a declared shape agrees with an + independent computation, not that the pass agrees with itself. + """ + operand_typespecs = [typespecs[value_id] for value_id in node.input_value_ids] + + if isinstance(node.operator, MatmulOperator): + assert len(operand_typespecs) == 2, "a matmul is a two-operand operation" + left, right = operand_typespecs + left_shape = list(left["shape"]) # type: ignore[index] + right_shape = list(right["shape"]) # type: ignore[index] + assert len(left_shape) == 2 and len(right_shape) == 2 + assert left_shape[1] == right_shape[0], "matmul inner dimensions must agree" + assert left["dtype"] == right["dtype"] # type: ignore[index] + return _typespec(str(left["dtype"]), (left_shape[0], right_shape[1])) # type: ignore[index] + + if isinstance(node.operator, MulOperator): + assert len(operand_typespecs) == 1, "an emitted mul scales by a literal" + operand = operand_typespecs[0] + return _typespec(str(operand["dtype"]), list(operand["shape"])) # type: ignore[index] + + if isinstance(node.operator, ReshapeOperator): + assert len(operand_typespecs) == 1 + operand = operand_typespecs[0] + source_elements = 1 + for dimension in operand["shape"]: # type: ignore[index] + source_elements *= int(dimension) + target_shape = list(node.op_params["shape"]) + target_elements = 1 + for dimension in target_shape: + target_elements *= int(dimension) + assert source_elements == target_elements, "a reshape preserves element count" + return _typespec(str(operand["dtype"]), target_shape) # type: ignore[index] + + # A fill node produces exactly what its descriptor declares. + from tinychain.autodiff import fill_descriptor + + descriptor = fill_descriptor(node) + assert not operand_typespecs, "a fill node has no operand" + return _typespec(descriptor.dtype, descriptor.shape) + + +# -------------------------------------------------------------------------- +# AC-1 — the rank-preserving tier emits exactly the five specified nodes +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize(("rows", "columns"), [(3, 5), (1, 1), (5, 3)]) +def test_rank_preserving_expansion_emits_the_five_specified_nodes( + rows: int, columns: int +) -> None: + from tinychain.autodiff import FillOperator + + graph = _traced_mean_graph(shape=(rows, columns), keepdims=True) + mean_node = graph.nodes[0] + operand_value_id = mean_node.input_value_ids[0] + + expanded = _expand(graph) + + assert len(expanded.nodes) == 5 + first_fill, row_sum, second_fill, total_sum, scale = expanded.nodes + + assert isinstance(first_fill.operator, FillOperator) + assert first_fill.input_value_ids == [] + assert first_fill.op_params == {"fill": 1.0, "dtype": "f64", "shape": [columns, 1]} + assert first_fill.output_typespec == _typespec("f64", (columns, 1)) + + assert isinstance(row_sum.operator, MatmulOperator) + assert row_sum.input_value_ids == [operand_value_id, first_fill.output_value_id] + assert row_sum.op_params == {} + assert row_sum.output_typespec == _typespec("f64", (rows, 1)) + + assert isinstance(second_fill.operator, FillOperator) + assert second_fill.input_value_ids == [] + assert second_fill.op_params == {"fill": 1.0, "dtype": "f64", "shape": [1, rows]} + assert second_fill.output_typespec == _typespec("f64", (1, rows)) + + assert isinstance(total_sum.operator, MatmulOperator) + assert total_sum.input_value_ids == [ + second_fill.output_value_id, + row_sum.output_value_id, + ] + assert total_sum.op_params == {} + assert total_sum.output_typespec == _typespec("f64", (1, 1)) + + assert isinstance(scale.operator, MulOperator) + assert scale.input_value_ids == [total_sum.output_value_id] + assert scale.op_params == {"right_literal": 1.0 / (rows * columns)} + assert scale.output_value_id == mean_node.output_value_id + assert scale.output_typespec == mean_node.output_typespec + + assert not any(isinstance(node.operator, MeanOperator) for node in expanded.nodes) + + +# -------------------------------------------------------------------------- +# AC-2 — the rank-reducing tier appends a real reshape to rank zero +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize(("rows", "columns"), [(3, 5), (1, 1)]) +def test_rank_reducing_expansion_appends_a_real_reshape_to_rank_zero( + rows: int, columns: int +) -> None: + graph = _traced_mean_graph(shape=(rows, columns), keepdims=False) + mean_node = graph.nodes[0] + assert mean_node.output_typespec == _typespec("f64", ()) + + expanded = _expand(graph) + + assert len(expanded.nodes) == 6 + scale = expanded.nodes[4] + reshape = expanded.nodes[5] + + assert isinstance(scale.operator, MulOperator) + assert scale.output_typespec == _typespec("f64", (1, 1)) + assert scale.output_value_id != mean_node.output_value_id + + assert isinstance(reshape.operator, ReshapeOperator) + assert reshape.input_value_ids == [scale.output_value_id] + assert reshape.op_params == {"shape": []} + assert reshape.output_value_id == mean_node.output_value_id + assert reshape.output_typespec == mean_node.output_typespec + assert reshape.output_typespec == _typespec("f64", ()) + + +# -------------------------------------------------------------------------- +# AC-3 — the truthful-shape audit, over both tiers +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize("keepdims", [True, False]) +@pytest.mark.parametrize(("rows", "columns"), [(3, 5), (1, 1), (5, 3)]) +def test_every_emitted_node_declares_the_shape_its_operation_produces( + keepdims: bool, rows: int, columns: int +) -> None: + graph = _traced_mean_graph(shape=(rows, columns), keepdims=keepdims) + + expanded = _expand(graph) + typespecs = _value_typespecs(expanded) + + for node in expanded.nodes: + assert node.output_typespec is not None, f"node {node.node_id!r} declares no typespec" + assert node.output_typespec == _independently_computed_typespec(node, typespecs), ( + f"node {node.node_id!r} declares a shape its operation does not produce" + ) + + +# -------------------------------------------------------------------------- +# AC-4 — the permitted operator set, and the two-operand matmul contract +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize("keepdims", [True, False]) +def test_every_emitted_node_is_in_the_permitted_operator_set(keepdims: bool) -> None: + from tinychain.autodiff import FillOperator + + graph = _traced_mean_graph(keepdims=keepdims) + + expanded = _expand(graph) + + for node in expanded.nodes: + assert isinstance( + node.operator, (FillOperator, MatmulOperator, MulOperator, ReshapeOperator) + ), f"node {node.node_id!r} emits {type(node.operator).__name__}" + if isinstance(node.operator, ReshapeOperator): + assert not keepdims, "a reshape is emitted only in the rank-reducing tier" + if isinstance(node.operator, MatmulOperator): + assert len(node.input_value_ids) == 2 + + +# -------------------------------------------------------------------------- +# AC-5 — boundaries and value ids are preserved +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize("keepdims", [True, False]) +def test_expansion_preserves_inputs_outputs_and_every_pre_existing_value_id( + keepdims: bool, +) -> None: + graph = _traced_mean_graph(keepdims=keepdims) + value_ids_before = _value_ids(graph) + + expanded = _expand(graph) + + assert expanded.inputs == graph.inputs + assert expanded.outputs == graph.outputs + assert value_ids_before <= _value_ids(expanded) + + +# -------------------------------------------------------------------------- +# AC-6 — purity and determinism +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize("keepdims", [True, False]) +def test_expansion_does_not_mutate_the_input_graph(keepdims: bool) -> None: + graph = _traced_mean_graph(keepdims=keepdims) + before = copy.deepcopy(graph) + + _expand(graph) + + assert graph == before + + +@pytest.mark.parametrize("keepdims", [True, False]) +def test_expansion_of_equal_graphs_returns_equal_graphs(keepdims: bool) -> None: + first = _traced_mean_graph(keepdims=keepdims) + second = _traced_mean_graph(keepdims=keepdims) + assert first == second + + assert _expand(first) == _expand(second) + assert _expand(first) == _expand(copy.deepcopy(first)) + + +# -------------------------------------------------------------------------- +# AC-7 — the eight-clause predicate and its categorized failures +# -------------------------------------------------------------------------- + +_FAILURE_CASES: list[tuple[str, dict[str, object], str, int]] = [ + ( + "two_operands", + {"input_value_ids": ["v0", "v0"]}, + "unsupported_reduction", + 1, + ), + ( + "operand_typespec_absent", + {"operand_typespec": None}, + "missing_shape_metadata", + 2, + ), + ( + "operand_typespec_without_shape", + {"operand_typespec": {"dtype": "f64"}}, + "missing_shape_metadata", + 2, + ), + ( + "operand_typespec_without_dtype", + {"operand_typespec": {"shape": [3, 5]}}, + "missing_dtype_metadata", + 2, + ), + ( + "operand_rank_one", + { + "operand_typespec": {"dtype": "f64", "shape": [5]}, + "op_params": {"axes": [0], "keepdims": True}, + "output_typespec": {"dtype": "f64", "shape": [1]}, + }, + "unsupported_reduction", + 2, + ), + ( + "operand_rank_three", + { + "operand_typespec": {"dtype": "f64", "shape": [2, 3, 5]}, + "op_params": {"axes": [0, 1, 2], "keepdims": True}, + "output_typespec": {"dtype": "f64", "shape": [1, 1, 1]}, + }, + "unsupported_reduction", + 2, + ), + ( + "symbolic_reduced_dimension", + {"operand_typespec": {"dtype": "f64", "shape": ["rows", 5]}}, + "unresolved_symbolic_shape", + 3, + ), + ( + "zero_reduced_dimension", + {"operand_typespec": {"dtype": "f64", "shape": [0, 5]}}, + "unsupported_reduction", + 3, + ), + ( + "partial_axes", + { + "op_params": {"axes": [0], "keepdims": True}, + "output_typespec": {"dtype": "f64", "shape": [1, 5]}, + }, + "unsupported_reduction", + 4, + ), + ( + "duplicated_axes", + {"op_params": {"axes": [0, 0], "keepdims": True}}, + "reduction_shape_mismatch", + 4, + ), + ( + "out_of_range_axis", + {"op_params": {"axes": [0, 2], "keepdims": True}}, + "reduction_shape_mismatch", + 4, + ), + ( + "malformed_axes", + {"op_params": {"axes": "both", "keepdims": True}}, + "reduction_shape_mismatch", + 4, + ), + ( + "missing_axes", + {"op_params": {"keepdims": True}}, + "reduction_shape_mismatch", + 4, + ), + ( + "keepdims_not_a_bool", + {"op_params": {"axes": [0, 1], "keepdims": 1}}, + "unsupported_reduction", + 5, + ), + ( + "missing_keepdims", + {"op_params": {"axes": [0, 1]}}, + "unsupported_reduction", + 5, + ), + ( + "non_floating_dtype", + { + "operand_dtype": "i32", + "operand_typespec": {"dtype": "i32", "shape": [3, 5]}, + "output_typespec": {"dtype": "i32", "shape": [1, 1]}, + }, + "dtype_not_differentiable", + 6, + ), + ( + "output_typespec_absent", + {"output_typespec": None}, + "missing_shape_metadata", + 7, + ), + ( + "output_typespec_without_dtype", + {"output_typespec": {"shape": [1, 1]}}, + "missing_dtype_metadata", + 7, + ), + ( + "output_shape_disagrees_with_the_reduction_rule", + { + "op_params": {"axes": [0, 1], "keepdims": False}, + "output_typespec": {"dtype": "f64", "shape": [1, 1]}, + }, + "reduction_shape_mismatch", + 7, + ), + ( + "output_dtype_disagrees_with_the_operand", + {"output_typespec": {"dtype": "f32", "shape": [1, 1]}}, + "reduction_shape_mismatch", + 7, + ), + ( + "descriptor_key_in_op_params", + {"op_params": {"axes": [0, 1], "keepdims": True, "fill": 1.0}}, + "unsupported_reduction", + 8, + ), +] + + +@pytest.mark.parametrize( + ("overrides", "expected_category", "clause"), + [case[1:] for case in _FAILURE_CASES], + ids=[case[0] for case in _FAILURE_CASES], +) +def test_an_unsupported_mean_raises_its_category_naming_the_node_and_the_clause( + overrides: dict[str, object], expected_category: str, clause: int +) -> None: + graph = _mean_graph(node_id="n_offender", **overrides) # type: ignore[arg-type] + + with pytest.raises(AutodiffError) as raised: + _expand(graph) + + assert raised.value.category == expected_category + assert "n_offender" in raised.value.message + assert f"clause {clause}" in raised.value.message + + +def test_a_mean_whose_node_id_is_reserved_fails_the_identifier_clause() -> None: + reserved = _reserved_node_id() + graph = _mean_graph(node_id=reserved) + + with pytest.raises(AutodiffError) as raised: + _expand(graph) + + assert raised.value.category == "unsupported_reduction" + assert reserved in raised.value.message + assert "clause 8" in raised.value.message + + +def test_a_mean_whose_output_value_id_is_reserved_fails_the_identifier_clause() -> None: + graph = _mean_graph(node_id="n_offender", output_value_id=_reserved_value_id()) + + with pytest.raises(AutodiffError) as raised: + _expand(graph) + + assert raised.value.category == "unsupported_reduction" + assert "n_offender" in raised.value.message + assert "clause 8" in raised.value.message + + +def test_a_rejected_mean_leaves_no_partially_rewritten_graph() -> None: + """Validation runs over every candidate before a single node is emitted.""" + supported = TensorNodeRecord( + node_id="n0", + output_value_id="v1", + operator=MeanOperator(), + op_params={"axes": [0, 1], "keepdims": True}, + input_value_ids=["v0"], + output_typespec=_typespec("f64", (1, 1)), + ) + unsupported = TensorNodeRecord( + node_id="n_offender", + output_value_id="v2", + operator=MeanOperator(), + op_params={"axes": [0], "keepdims": True}, + input_value_ids=["v0"], + output_typespec=_typespec("f64", (1, 5)), + ) + graph = TensorGraph( + nodes=[supported, unsupported], + inputs=[("v0", _typespec("f64", (3, 5)))], + outputs=["v1", "v2"], + ) + before = copy.deepcopy(graph) + + with pytest.raises(AutodiffError) as raised: + _expand(graph) + + assert raised.value.category == "unsupported_reduction" + assert "n_offender" in raised.value.message + assert graph == before + + +# -------------------------------------------------------------------------- +# AC-8 — unrelated nodes are carried through identically and in order +# -------------------------------------------------------------------------- + + +def _mixed_graph() -> TensorGraph: + """A graph whose mean is surrounded by nodes the pass must not touch.""" + column_sum = TensorNodeRecord( + node_id="n0", + output_value_id="v1", + operator=SumOperator(), + op_params={"axes": [0], "keepdims": True}, + input_value_ids=["v0"], + output_typespec=_typespec("f64", (1, 5)), + ) + mean = TensorNodeRecord( + node_id="n1", + output_value_id="v2", + operator=MeanOperator(), + op_params={"axes": [0, 1], "keepdims": True}, + input_value_ids=["v0"], + output_typespec=_typespec("f64", (1, 1)), + ) + reshape = TensorNodeRecord( + node_id="n2", + output_value_id="v3", + operator=ReshapeOperator(), + op_params={"shape": [5, 1]}, + input_value_ids=["v1"], + output_typespec=_typespec("f64", (5, 1)), + ) + return TensorGraph( + nodes=[column_sum, mean, reshape], + inputs=[("v0", _typespec("f64", (3, 5)))], + outputs=["v1", "v2", "v3"], + ) + + +def test_unmatched_nodes_are_carried_through_identically_and_in_order() -> None: + graph = _mixed_graph() + column_sum, mean, reshape = graph.nodes + + expanded = _expand(graph) + + assert len(expanded.nodes) == 7 + assert expanded.nodes[0] == column_sum + assert expanded.nodes[6] == reshape + assert [node.node_id for node in expanded.nodes].index("n0") < [ + node.node_id for node in expanded.nodes + ].index("n2") + assert not any(isinstance(node.operator, MeanOperator) for node in expanded.nodes) + assert mean.output_value_id == expanded.nodes[5].output_value_id + + +# -------------------------------------------------------------------------- +# AC-9 — reserved namespace, collisions, and duplicate identifiers +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize("keepdims", [True, False]) +def test_every_minted_identifier_comes_from_the_reserved_namespace(keepdims: bool) -> None: + from tinychain.autodiff.expansion import ( + EXPANSION_NODE_ID_PREFIX, + EXPANSION_VALUE_ID_PREFIX, + ) + + graph = _traced_mean_graph(keepdims=keepdims) + mean_node = graph.nodes[0] + + expanded = _expand(graph) + + node_ids = [node.node_id for node in expanded.nodes] + assert len(set(node_ids)) == len(node_ids) + assert all(node_id.startswith(EXPANSION_NODE_ID_PREFIX) for node_id in node_ids) + + minted_value_ids = [ + node.output_value_id + for node in expanded.nodes + if node.output_value_id != mean_node.output_value_id + ] + assert len(set(minted_value_ids)) == len(minted_value_ids) + assert all( + value_id.startswith(EXPANSION_VALUE_ID_PREFIX) for value_id in minted_value_ids + ) + # Disjoint from the tracer's `v…`/`n…` and the reverse transform's `d…`/`dn…`. + for identifier in node_ids + minted_value_ids: + assert not identifier.startswith(("v", "n", "d")) + + +def test_a_minted_node_id_colliding_with_an_existing_one_is_rejected() -> None: + reserved = _reserved_node_id() + graph = _traced_mean_graph(keepdims=True) + squatter = TensorNodeRecord( + node_id=reserved, + output_value_id="v9", + operator=SumOperator(), + op_params={"axes": [0], "keepdims": True}, + input_value_ids=["v0"], + output_typespec=_typespec("f64", (1, 5)), + ) + seeded = TensorGraph( + nodes=[squatter, *graph.nodes], + inputs=list(graph.inputs), + outputs=[*graph.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_a_minted_value_id_colliding_with_an_existing_one_is_rejected() -> None: + reserved = _reserved_value_id() + graph = _traced_mean_graph(keepdims=True) + squatter = TensorNodeRecord( + node_id="n9", + output_value_id=reserved, + operator=SumOperator(), + op_params={"axes": [0], "keepdims": True}, + input_value_ids=["v0"], + output_typespec=_typespec("f64", (1, 5)), + ) + seeded = TensorGraph( + nodes=[squatter, *graph.nodes], + inputs=list(graph.inputs), + outputs=[*graph.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_an_input_graph_with_duplicate_node_ids_is_rejected() -> None: + graph = _traced_mean_graph(keepdims=True) + duplicate = TensorNodeRecord( + node_id=graph.nodes[0].node_id, + output_value_id="v9", + operator=SumOperator(), + op_params={"axes": [0], "keepdims": True}, + input_value_ids=["v0"], + output_typespec=_typespec("f64", (1, 5)), + ) + seeded = TensorGraph( + nodes=[*graph.nodes, duplicate], + inputs=list(graph.inputs), + outputs=[*graph.outputs, "v9"], + ) + + with pytest.raises(AutodiffError) as raised: + _expand(seeded) + + assert raised.value.category == "malformed_derivative_ir" + assert graph.nodes[0].node_id in raised.value.message + + +def test_an_input_graph_with_duplicate_value_ids_is_rejected() -> None: + graph = _traced_mean_graph(keepdims=True) + duplicate = TensorNodeRecord( + node_id="n9", + output_value_id=graph.nodes[0].output_value_id, + operator=SumOperator(), + op_params={"axes": [0], "keepdims": True}, + input_value_ids=["v0"], + output_typespec=_typespec("f64", (1, 5)), + ) + seeded = TensorGraph( + nodes=[*graph.nodes, duplicate], + inputs=list(graph.inputs), + outputs=list(graph.outputs), + ) + + with pytest.raises(AutodiffError) as raised: + _expand(seeded) + + assert raised.value.category == "malformed_derivative_ir" + assert graph.nodes[0].output_value_id in raised.value.message + + +# -------------------------------------------------------------------------- +# export surface +# -------------------------------------------------------------------------- + + +def test_expand_mean_graph_is_exported_from_the_autodiff_package() -> None: + from tinychain import autodiff + + assert "expand_mean_graph" in autodiff.__all__ + assert callable(autodiff.expand_mean_graph) + assert not hasattr(tc, "expand_mean_graph") diff --git a/py/tinychain/autodiff/__init__.py b/py/tinychain/autodiff/__init__.py index 062352f..62edfc0 100644 --- a/py/tinychain/autodiff/__init__.py +++ b/py/tinychain/autodiff/__init__.py @@ -119,6 +119,7 @@ { "FillDescriptor", "FillOperator", + "expand_mean_graph", "fill_descriptor", } ) @@ -292,6 +293,7 @@ def __getattr__(name: str) -> object: "analyze_graph_dependencies", "FillDescriptor", "FillOperator", + "expand_mean_graph", "fill_descriptor", "LOWERING_CLAIM_HANDLER", "LOWERING_CLAIM_FUSION", diff --git a/py/tinychain/autodiff/expansion.py b/py/tinychain/autodiff/expansion.py index 2695684..ff16647 100644 --- a/py/tinychain/autodiff/expansion.py +++ b/py/tinychain/autodiff/expansion.py @@ -64,6 +64,51 @@ Unknown keys are checked before the individual fields, and ``fill``, ``shape``, and ``dtype`` are then checked in that order, so an absent ``shape`` reports ``missing_shape_metadata`` rather than being swallowed by the key-set check. + +The forward mean expansion +-------------------------- +:func:`expand_mean_graph` rewrites every *supported* all-axis rank-2 +``MeanOperator`` in a :class:`~tinychain.autodiff.graph.TensorGraph` into a +region built from a generated constant, two matmuls, and one scale -- so a +backend needs no reduction operation at all. For an operand of shape ``(r, c)`` +and dtype ``D``:: + + f1 FillOperator -- -> D, [c, 1] ones + f2 MatmulOperator [operand, f1] -> D, [r, 1] row sums + f3 FillOperator -- -> D, [1, r] ones + f4 MatmulOperator [f3, f2] -> D, [1, 1] total sum + f5 MulOperator [f4] -> D, [1, 1] mean, rank 2 + +Two tiers follow from the mean's own ``keepdims``, with no parameter. For +``keepdims=True`` the mean is already ``[1, 1]``, so ``f5`` carries the mean's +value id and typespec and the region is five nodes. For ``keepdims=False`` the +mean is rank zero, so a sixth node -- a real ``ReshapeOperator`` with +``shape = []`` -- performs the genuine rank change and carries the mean's value +id. A rank change is never expressed by an elementwise node declaring a +different shape: a scalar and a ``[1, 1]`` value are not interchangeable, and +every emitted node's ``output_typespec`` is computed from its operands rather +than copied from the node it replaces. + +An unsupported mean stops the pass rather than being left in place -- a mean is +unambiguously inside the declared domain of a pass named for mean expansion. +Every candidate is validated before any node is emitted, so a rejected graph +never comes back partially rewritten, and every failure names the offending node +and the clause of FR-128-006 it failed. Operators other than ``MeanOperator`` +are never rewritten and are carried through identical. + +The reserved identifier namespace +--------------------------------- +Nodes and values a pass creates are named ``exn0, exn1, …`` and +``exv0, exv1, …`` -- see :data:`EXPANSION_NODE_ID_PREFIX` and +:data:`EXPANSION_VALUE_ID_PREFIX`. The namespace is reserved for expansion and +is disjoint from the tracer's ``v…``/``n…`` and the reverse transform's +``d…``/``dn…``, so a minted identifier can never be mistaken for a traced or +generated one. Minting is deterministic -- indices run from zero on every call, +so equal graphs expand to equal graphs including identifiers and order -- and +every minted identifier is checked against those already in the artifact, +raising ``malformed_derivative_ir`` naming the identifier rather than silently +shadowing an existing value. An artifact that already contains a duplicate node +or value id is rejected the same way, before anything is minted. """ from __future__ import annotations @@ -71,9 +116,27 @@ from collections.abc import Mapping, Sequence from dataclasses import dataclass -from .graph import TensorNodeRecord, TensorOperator +from .graph import ( + MatmulOperator, + MeanOperator, + MulOperator, + ReshapeOperator, + TensorGraph, + TensorNodeRecord, + TensorOperator, +) from .lowering import OperationContext from .protocol import AutodiffError +from .shape import ( + Shape, + _normalize_mean_axes, + check_compatible_operand_dtypes, + check_differentiable_dtype, + matmul_output_shape, + mean_output_shape, + shape_rank, + typespec_ranked_shape, +) # The route name a fill node reports. It is not a dispatch key -- handlers are # selected by concrete operator type -- and exists for serialization parity with @@ -215,3 +278,564 @@ def _read_shape(node_id: str, op_params: Mapping[str, object]) -> tuple[int, ... ) dimensions.append(dimension) return tuple(dimensions) + + +# -------------------------------------------------------------------------- +# The reserved identifier namespace (§9.2) +# -------------------------------------------------------------------------- + +# Every node and value an expansion pass creates is named from this reserved +# namespace. It is disjoint from the tracer's `v…`/`n…` and from the reverse +# transform's `d…`/`dn…`, so a minted identifier can never be confused with a +# traced or generated one; minting is deterministic (indices run from zero on +# every call, so equal artifacts expand to equal artifacts); and every minted +# identifier is checked against the identifiers already in the artifact, which +# fails closed rather than silently shadowing an existing value. +EXPANSION_NODE_ID_PREFIX = "exn" +EXPANSION_VALUE_ID_PREFIX = "exv" + + +class _IdentifierMinter: + """Deterministic source of reserved node and value identifiers. + + Constructed once per pass invocation over the identifiers the input + artifact already uses. Both mint methods take no argument, so every caller + -- this pass and the gradient-path pass that shares this module -- mints + from one place and cannot disagree about the namespace or the ordering. + """ + + def __init__( + self, *, existing_node_ids: frozenset[str], existing_value_ids: frozenset[str] + ) -> None: + self._existing_node_ids = existing_node_ids + self._existing_value_ids = existing_value_ids + self._node_index = 0 + self._value_index = 0 + + def mint_node_id(self) -> str: + """Return the next reserved node id, rejecting a collision.""" + node_id = f"{EXPANSION_NODE_ID_PREFIX}{self._node_index}" + self._node_index += 1 + if node_id in self._existing_node_ids: + raise AutodiffError( + "malformed_derivative_ir", + f"minted node id {node_id!r} collides with a node id already present in " + f"the artifact; {EXPANSION_NODE_ID_PREFIX!r} is a reserved expansion namespace", + ) + return node_id + + def mint_value_id(self) -> str: + """Return the next reserved value id, rejecting a collision.""" + value_id = f"{EXPANSION_VALUE_ID_PREFIX}{self._value_index}" + self._value_index += 1 + if value_id in self._existing_value_ids: + raise AutodiffError( + "malformed_derivative_ir", + f"minted value id {value_id!r} collides with a value id already present in " + f"the artifact; {EXPANSION_VALUE_ID_PREFIX!r} is a reserved expansion namespace", + ) + return value_id + + +def _is_reserved_identifier(identifier: object) -> bool: + """Report whether *identifier* is spelled inside the reserved namespace.""" + return isinstance(identifier, str) and identifier.startswith( + (EXPANSION_NODE_ID_PREFIX, EXPANSION_VALUE_ID_PREFIX) + ) + + +# -------------------------------------------------------------------------- +# Artifact identifier indexing +# -------------------------------------------------------------------------- + + +def _indexed_node_ids(nodes: Sequence[TensorNodeRecord]) -> frozenset[str]: + """Return every node id in *nodes*, rejecting a duplicate.""" + seen: set[str] = set() + for node in nodes: + if node.node_id in seen: + raise AutodiffError( + "malformed_derivative_ir", + f"artifact declares duplicate node id {node.node_id!r}", + ) + seen.add(node.node_id) + return frozenset(seen) + + +def _indexed_value_ids( + nodes: Sequence[TensorNodeRecord], inputs: Sequence[tuple[str, object]] +) -> frozenset[str]: + """Return every declared value id, rejecting a duplicate producer.""" + seen: set[str] = set() + for value_id, _ in inputs: + if value_id in seen: + raise AutodiffError( + "malformed_derivative_ir", + f"artifact declares duplicate value id {value_id!r}", + ) + seen.add(value_id) + for node in nodes: + if node.output_value_id in seen: + raise AutodiffError( + "malformed_derivative_ir", + f"artifact declares duplicate value id {node.output_value_id!r}", + ) + seen.add(node.output_value_id) + return frozenset(seen) + + +def _declared_typespecs( + nodes: Sequence[TensorNodeRecord], inputs: Sequence[tuple[str, object]] +) -> dict[str, object]: + """Index every value id to the typespec its producer declares for it.""" + typespecs: dict[str, object] = {value_id: typespec for value_id, typespec in inputs} + for node in nodes: + typespecs[node.output_value_id] = node.output_typespec + return typespecs + + +# -------------------------------------------------------------------------- +# The supported-mean predicate (FR-128-006) +# -------------------------------------------------------------------------- + +# The rank the forward expansion is defined for. A matmul pair reduces exactly +# two axes, so a mean over any other rank is out of the pass's declared domain +# rather than a metadata defect. +_SUPPORTED_MEAN_RANK = 2 + + +@dataclass(frozen=True) +class _SupportedMean: + """One `MeanOperator` node proven expandable, with everything the emitter needs.""" + + node_id: str + operand_value_id: str + operand_typespec: dict[str, object] + dtype: str + rows: int + columns: int + keepdims: bool + output_value_id: str + output_typespec: dict[str, object] + + +def _mean_failure(node_id: str, clause: int, category: str, detail: str) -> AutodiffError: + """Build the categorized failure for one clause of FR-128-006. + + Every message names both the offending node and the clause that failed, so + a caller reading only the message can find the node and the rule. + """ + return AutodiffError( + category, + f"mean node {node_id!r} fails clause {clause} of FR-128-006: {detail}", + ) + + +def _recategorized( + node_id: str, clause: int, detail: str, exc: AutodiffError +) -> AutodiffError: + """Re-raise a helper's own categorized failure with node and clause context.""" + return _mean_failure(node_id, clause, exc.category, f"{detail}: {exc.message}") + + +def _mean_operand_shape_and_dtype( + node: TensorNodeRecord, typespecs: Mapping[str, object] +) -> tuple[Shape, str]: + """Clause 2 -- the operand carries a complete typespec of rank two.""" + operand_typespec = typespecs.get(node.input_value_ids[0]) + try: + shape = typespec_ranked_shape(operand_typespec) + except AutodiffError as exc: + raise _recategorized( + node.node_id, 2, "operand declares no ranked shape", exc + ) from exc + + dtype = None if operand_typespec is None else operand_typespec.get("dtype") + if not isinstance(dtype, str) or not dtype: + raise _mean_failure( + node.node_id, + 2, + "missing_dtype_metadata", + f"operand declares no dtype: {dtype!r}", + ) + + if shape_rank(shape) != _SUPPORTED_MEAN_RANK: + raise _mean_failure( + node.node_id, + 2, + "unsupported_reduction", + f"operand shape {list(shape)!r} has rank {shape_rank(shape)}, and this pass " + f"expands a mean over a rank-{_SUPPORTED_MEAN_RANK} operand only", + ) + return shape, dtype + + +def _mean_operand_dimensions(node_id: str, shape: Shape) -> tuple[int, int]: + """Clause 3 -- both operand dimensions are positive integers.""" + for axis, dimension in enumerate(shape): + if isinstance(dimension, str): + raise _mean_failure( + node_id, + 3, + "unresolved_symbolic_shape", + f"reduced dimension {dimension!r} at axis {axis} is symbolic, not an integer", + ) + if dimension <= 0: + raise _mean_failure( + node_id, + 3, + "unsupported_reduction", + f"reduced dimension {dimension!r} at axis {axis} is not positive", + ) + rows, columns = shape + return int(rows), int(columns) + + +def _mean_axes(node: TensorNodeRecord) -> tuple[int, ...]: + """Clause 4 -- the declared axes normalize to every axis of the operand.""" + axes = node.op_params.get("axes") + try: + normalized = _normalize_mean_axes(axes, _SUPPORTED_MEAN_RANK) + except AutodiffError as exc: + raise _recategorized(node.node_id, 4, "declared axes are malformed", exc) from exc + if set(normalized) != set(range(_SUPPORTED_MEAN_RANK)): + raise _mean_failure( + node.node_id, + 4, + "unsupported_reduction", + f"declared axes {list(normalized)!r} are a partial reduction, and this pass " + "expands an all-axis mean only", + ) + return normalized + + +def _mean_keepdims(node: TensorNodeRecord) -> bool: + """Clause 5 -- ``keepdims`` is a real boolean; both values select a tier.""" + keepdims = node.op_params.get("keepdims") + if not isinstance(keepdims, bool): + raise _mean_failure( + node.node_id, + 5, + "unsupported_reduction", + f"declared 'keepdims' is not a bool: {keepdims!r}", + ) + return keepdims + + +def _mean_dtype(node_id: str, dtype: str) -> str: + """Clause 6 -- the operand dtype is differentiable.""" + try: + return check_differentiable_dtype(dtype) + except AutodiffError as exc: + raise _recategorized(node_id, 6, "operand dtype is not differentiable", exc) from exc + + +def _mean_output_typespec( + node: TensorNodeRecord, + operand_shape: Shape, + operand_dtype: str, + axes: tuple[int, ...], + keepdims: bool, +) -> dict[str, object]: + """Clause 7 -- the declared output typespec is complete and is the true one.""" + declared = node.output_typespec + try: + declared_shape = typespec_ranked_shape(declared) + except AutodiffError as exc: + raise _recategorized( + node.node_id, 7, "output declares no ranked shape", exc + ) from exc + + declared_dtype = None if declared is None else declared.get("dtype") + if not isinstance(declared_dtype, str) or not declared_dtype: + raise _mean_failure( + node.node_id, + 7, + "missing_dtype_metadata", + f"output declares no dtype: {declared_dtype!r}", + ) + + expected_shape = mean_output_shape(operand_shape, list(axes), keepdims=keepdims) + if tuple(declared_shape) != tuple(expected_shape): + raise _mean_failure( + node.node_id, + 7, + "reduction_shape_mismatch", + f"output declares shape {list(declared_shape)!r}, but reducing " + f"{list(operand_shape)!r} over {list(axes)!r} with keepdims={keepdims} " + f"produces {list(expected_shape)!r}", + ) + if declared_dtype != operand_dtype: + raise _mean_failure( + node.node_id, + 7, + "reduction_shape_mismatch", + f"output declares dtype {declared_dtype!r}, but the operand dtype is " + f"{operand_dtype!r} and a mean performs no promotion", + ) + return {"dtype": declared_dtype, "shape": list(declared_shape)} + + +def _check_mean_carries_no_reserved_construct(node: TensorNodeRecord) -> None: + """Clause 8 -- the node carries no reserved identifier and no fill descriptor key.""" + for label, identifier in ( + ("node id", node.node_id), + ("output value id", node.output_value_id), + ): + if _is_reserved_identifier(identifier): + raise _mean_failure( + node.node_id, + 8, + "unsupported_reduction", + f"{label} {identifier!r} is spelled inside the reserved expansion namespace " + f"({EXPANSION_NODE_ID_PREFIX!r}/{EXPANSION_VALUE_ID_PREFIX!r})", + ) + descriptor_keys = sorted(str(key) for key in node.op_params if key in _FILL_PARAM_KEYS) + if descriptor_keys: + raise _mean_failure( + node.node_id, + 8, + "unsupported_reduction", + f"op_params carries fill descriptor key(s) {descriptor_keys!r}", + ) + + +def _supported_mean( + node: TensorNodeRecord, typespecs: Mapping[str, object] +) -> _SupportedMean: + """Prove one `MeanOperator` node expandable, or raise its categorized failure. + + The eight clauses of FR-128-006 are checked in their stated order, so the + reported clause is always the first one the node fails. + """ + if len(node.input_value_ids) != 1: + raise _mean_failure( + node.node_id, + 1, + "unsupported_reduction", + f"a mean takes exactly one operand, got {list(node.input_value_ids)!r}", + ) + + operand_shape, operand_dtype = _mean_operand_shape_and_dtype(node, typespecs) + rows, columns = _mean_operand_dimensions(node.node_id, operand_shape) + axes = _mean_axes(node) + keepdims = _mean_keepdims(node) + dtype = _mean_dtype(node.node_id, operand_dtype) + output_typespec = _mean_output_typespec( + node, operand_shape, operand_dtype, axes, keepdims + ) + _check_mean_carries_no_reserved_construct(node) + + return _SupportedMean( + node_id=node.node_id, + operand_value_id=node.input_value_ids[0], + operand_typespec={"dtype": dtype, "shape": [rows, columns]}, + dtype=dtype, + rows=rows, + columns=columns, + keepdims=keepdims, + output_value_id=node.output_value_id, + output_typespec=output_typespec, + ) + + +# -------------------------------------------------------------------------- +# The §8.3 region emitter +# -------------------------------------------------------------------------- + +# The value every generated ones-tensor holds. A row of ones is what turns a +# matmul into a sum along one axis, which is the whole mechanism of the region. +_ONES_FILL = 1.0 + + +def _boundary_typespec(dtype: str, shape: Sequence[object]) -> dict[str, object]: + """Build the graph-boundary typespec form, `{"dtype": str, "shape": list}`.""" + return {"dtype": dtype, "shape": list(shape)} + + +def _emit_fill( + minter: _IdentifierMinter, *, dtype: str, shape: Sequence[int] +) -> TensorNodeRecord: + """Emit one ones-tensor of *shape*; its true result is its own descriptor.""" + descriptor_shape = [int(dimension) for dimension in shape] + return TensorNodeRecord( + node_id=minter.mint_node_id(), + output_value_id=minter.mint_value_id(), + operator=FillOperator(), + op_params={"fill": _ONES_FILL, "dtype": dtype, "shape": descriptor_shape}, + input_value_ids=[], + output_typespec=_boundary_typespec(dtype, descriptor_shape), + ) + + +def _emit_matmul( + minter: _IdentifierMinter, + *, + left: TensorNodeRecord | tuple[str, Mapping[str, object]], + right: TensorNodeRecord | tuple[str, Mapping[str, object]], +) -> TensorNodeRecord: + """Emit one two-operand matmul, deriving its declared type from its operands.""" + left_value_id, left_typespec = _operand(left) + right_value_id, right_typespec = _operand(right) + dtype = check_compatible_operand_dtypes( + str(left_typespec["dtype"]), str(right_typespec["dtype"]) + ) + shape = matmul_output_shape( + tuple(left_typespec["shape"]), + tuple(right_typespec["shape"]), + ) + return TensorNodeRecord( + node_id=minter.mint_node_id(), + output_value_id=minter.mint_value_id(), + operator=MatmulOperator(), + op_params={}, + input_value_ids=[left_value_id, right_value_id], + output_typespec=_boundary_typespec(dtype, shape), + ) + + +def _emit_scale( + minter: _IdentifierMinter, + *, + operand: TensorNodeRecord, + right_literal: float, + output_value_id: str | None = None, +) -> TensorNodeRecord: + """Emit one scale-by-literal; an elementwise scale keeps the operand's type.""" + operand_value_id, operand_typespec = _operand(operand) + return TensorNodeRecord( + node_id=minter.mint_node_id(), + output_value_id=( + minter.mint_value_id() if output_value_id is None else output_value_id + ), + operator=MulOperator(), + op_params={"right_literal": float(right_literal)}, + input_value_ids=[operand_value_id], + output_typespec=_boundary_typespec( + str(operand_typespec["dtype"]), list(operand_typespec["shape"]) + ), + ) + + +def _emit_reshape( + minter: _IdentifierMinter, + *, + operand: TensorNodeRecord, + shape: Sequence[int], + output_value_id: str, +) -> TensorNodeRecord: + """Emit one real reshape; a reshape keeps the dtype and takes the target shape.""" + operand_value_id, operand_typespec = _operand(operand) + target_shape = [int(dimension) for dimension in shape] + return TensorNodeRecord( + node_id=minter.mint_node_id(), + output_value_id=output_value_id, + operator=ReshapeOperator(), + op_params={"shape": target_shape}, + input_value_ids=[operand_value_id], + output_typespec=_boundary_typespec(str(operand_typespec["dtype"]), target_shape), + ) + + +def _operand( + source: TensorNodeRecord | tuple[str, Mapping[str, object]], +) -> tuple[str, Mapping[str, object]]: + """Normalize an emitted node or an existing `(value id, typespec)` pair.""" + if isinstance(source, TensorNodeRecord): + return source.output_value_id, source.output_typespec + return source + + +def _emit_mean_region( + supported: _SupportedMean, minter: _IdentifierMinter +) -> list[TensorNodeRecord]: + """Emit the region of §8.3 replacing one proven-supported mean. + + Every declared shape below is computed from the operands the node actually + reads -- the shape helpers derive it -- rather than copied from the mean the + region replaces, so no node can declare a shape its operation cannot + produce. + """ + rows, columns = supported.rows, supported.columns + + column_ones = _emit_fill(minter, dtype=supported.dtype, shape=(columns, 1)) + row_sums = _emit_matmul( + minter, + left=(supported.operand_value_id, supported.operand_typespec), + right=column_ones, + ) + 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) + 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. + scale = _emit_scale( + minter, + operand=total_sum, + right_literal=reciprocal_count, + output_value_id=supported.output_value_id, + ) + return [column_ones, row_sums, row_ones, total_sum, scale] + + # The rank-reducing tier: the scale is `[1, 1]` and the mean is rank zero, + # so a real reshape performs the rank change and carries the mean's value id. + scale = _emit_scale(minter, operand=total_sum, right_literal=reciprocal_count) + rank_change = _emit_reshape( + minter, + operand=scale, + shape=(), + output_value_id=supported.output_value_id, + ) + return [column_ones, row_sums, row_ones, total_sum, scale, rank_change] + + +# -------------------------------------------------------------------------- +# The public forward pass +# -------------------------------------------------------------------------- + + +def expand_mean_graph(graph: TensorGraph) -> TensorGraph: + """Return *graph* with every supported all-axis rank-2 mean expanded. + + Each supported `MeanOperator` (FR-128-006) is replaced in place by the + matmul-based region of §8.3 -- five nodes when the mean declared + ``keepdims=True``, six when it declared ``keepdims=False``, the sixth a real + `ReshapeOperator` performing the rank change. Every other node, and + ``inputs`` and ``outputs``, are carried through unchanged; the input graph is + never mutated, and equal graphs expand to equal graphs. + + Every candidate mean is validated before a single node is emitted, so an + unsupported mean raises its categorized failure (§13.1) naming the node and + the clause of FR-128-006 it failed, and never yields a partially rewritten + graph. + """ + nodes = graph.nodes + existing_node_ids = _indexed_node_ids(nodes) + existing_value_ids = _indexed_value_ids(nodes, graph.inputs) + typespecs = _declared_typespecs(nodes, graph.inputs) + + # Validation first, over every candidate, before any emission (§13.2). + supported_means = { + node.node_id: _supported_mean(node, typespecs) + for node in nodes + if isinstance(node.operator, MeanOperator) + } + + minter = _IdentifierMinter( + existing_node_ids=existing_node_ids, existing_value_ids=existing_value_ids + ) + expanded_nodes: list[TensorNodeRecord] = [] + for node in nodes: + supported = supported_means.get(node.node_id) + if supported is None: + expanded_nodes.append(node) + continue + expanded_nodes.extend(_emit_mean_region(supported, minter)) + + return TensorGraph( + nodes=expanded_nodes, + inputs=list(graph.inputs), + outputs=list(graph.outputs), + )