fix(autodiff): close four gaps between the two expansion passes

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:

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