diff --git a/py/tests/test_autodiff_elementwise.py b/py/tests/test_autodiff_elementwise.py index 21c3b03..4fd6c37 100644 --- a/py/tests/test_autodiff_elementwise.py +++ b/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") diff --git a/py/tinychain/autodiff/vjp.py b/py/tinychain/autodiff/vjp.py index 0ec2290..dfb4374 100644 --- a/py/tinychain/autodiff/vjp.py +++ b/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)