fix(autodiff): accumulate repeated multiplication inputs

py/tests/test_autodiff_elementwise.py

@@ -4,6 +4,7 @@
 import pytest
 
 from tinychain.autodiff import (
+    AddOperator,
     AutodiffError,
     BroadcastReduceOperator,
     DerivativeMetadata,
@@ -111,6 +112,40 @@ def test_sub_vjp_broadcast_rhs_reduces_negated_gradient():
     np.testing.assert_allclose(dy, -np.sum(seed, axis=0, keepdims=True), rtol=1e-5)
 
 
+def test_mul_vjp_repeated_input_accumulates_both_partials():
+    graph = TensorGraph(
+        nodes=[
+            TensorNodeRecord(
+                node_id="n0",
+                output_value_id="v1",
+                operator=MulOperator(),
+                op_params={},
+                input_value_ids=["v0", "v0"],
+                output_typespec=_typespec((2, 3)),
+            )
+        ],
+        inputs=[("v0", _typespec((2, 3)))],
+        outputs=["v1"],
+    )
+    program = generate(graph, "v1", ["v0"], "seed")
+
+    assert _operator_types(program.nodes) == [MulOperator, MulOperator, AddOperator]
+    lhs_partial, rhs_partial, accumulated = program.nodes
+    assert accumulated.input_value_ids == [lhs_partial.output_value_id, rhs_partial.output_value_id]
+    assert program.gradients == {"v0": accumulated.output_value_id}
+    assert program.output_gradients == [accumulated.output_value_id]
+
+    seed = np.array([[1.0, 1.5, 2.0], [2.5, 3.0, 3.5]], dtype=np.float32)
+    value = np.array([[2.0, 4.0, 6.0], [8.0, 10.0, 12.0]], dtype=np.float32)
+    result = ExecutionScheduler(NumpyAutodiffDispatcher()).execute(
+        program,
+        values={"seed": seed, "v0": value},
+    )
+
+    (gradient,) = result.gradients
+    np.testing.assert_allclose(gradient, 2 * seed * value, rtol=1e-5)
+
+
 def test_mul_vjp_broadcast_rhs_executes_correct_gradients():
     graph = _elementwise_graph(MulOperator(), lhs_shape=(2, 3), rhs_shape=(1, 3), out_shape=(2, 3))
     program = generate(graph, "v2", ["v0", "v1"], "seed")

py/tinychain/autodiff/vjp.py

@@ -316,6 +316,8 @@ def apply(self, context: VjpContext) -> VjpResult:
         lhs_id, rhs_id, result_shape, result_typespec = self._validate_binary(context, "mul")
         gradients: dict[str, str] = {}
         derivative_nodes: list[TensorNodeRecord] = []
+        lhs_gradient_id: str | None = None
+        rhs_gradient_id: str | None = None
 
         if lhs_id in context.needed_input_value_ids:
             lhs_raw = _elementwise_binary_node(
@@ -325,13 +327,14 @@ def apply(self, context: VjpContext) -> VjpResult:
                 output_typespec=result_typespec,
             )
             derivative_nodes.append(lhs_raw)
-            gradients[lhs_id] = self._reduce_to_operand(
+            lhs_gradient_id = self._reduce_to_operand(
                 context=context,
                 gradient_id=lhs_raw.output_value_id,
                 operand_id=lhs_id,
                 result_shape=result_shape,
                 derivative_nodes=derivative_nodes,
             )
+            gradients[lhs_id] = lhs_gradient_id
 
         if rhs_id in context.needed_input_value_ids:
             rhs_raw = _elementwise_binary_node(
@@ -341,13 +344,24 @@ def apply(self, context: VjpContext) -> VjpResult:
                 output_typespec=result_typespec,
             )
             derivative_nodes.append(rhs_raw)
-            gradients[rhs_id] = self._reduce_to_operand(
+            rhs_gradient_id = self._reduce_to_operand(
                 context=context,
                 gradient_id=rhs_raw.output_value_id,
                 operand_id=rhs_id,
                 result_shape=result_shape,
                 derivative_nodes=derivative_nodes,
             )
+            gradients[rhs_id] = rhs_gradient_id
+
+        if lhs_id == rhs_id and lhs_gradient_id is not None and rhs_gradient_id is not None:
+            accumulated = _elementwise_binary_node(
+                context=context,
+                operator=AddOperator(),
+                input_value_ids=[lhs_gradient_id, rhs_gradient_id],
+                output_typespec=context.value_typespecs.get(lhs_id),
+            )
+            derivative_nodes.append(accumulated)
+            gradients[lhs_id] = accumulated.output_value_id
 
         return VjpResult(gradients=gradients, derivative_nodes=derivative_nodes)