Spaces:
Running
Running
| """Contract, numerical, and snapshot tests for the shared computation store.""" | |
| import json | |
| import numpy as np | |
| import pytest | |
| import sympy as sp | |
| from autodiff_visualizer import ComputationGraph, safe_parse_expr | |
| from backward_rules import ScalarLocalRule, TensorBackwardRule | |
| from diagram_renderer import DIVIDER, render_diagram, step_label | |
| from math_renderer import MathRenderer | |
| from tensor_engine import TensorComputationGraph | |
| from trace_renderer import _contribution_box_html, render_trace, render_trace_html | |
| def scalar(expr, feed): | |
| graph = ComputationGraph().build(safe_parse_expr(expr)) | |
| graph.evaluate({sp.Symbol(k): v for k, v in feed.items()}) | |
| return graph.to_computation_model() | |
| def tensor_graph(expr, feed, forward_only=False): | |
| graph = TensorComputationGraph().build_from_ast(expr) | |
| graph.evaluate(feed, forward_only=forward_only) | |
| return graph | |
| CASES = [ | |
| ('scalar', '(x*x)*y+(y+2)', {'x': 3., 'y': 4.}), | |
| ('scalar', 'sin(x) + x**y', {'x': 2., 'y': 3.}), | |
| ('tensor', 'sum(x/y+b)', {'x': np.array([2., 4.]), 'y': np.array([1., 2.]), 'b': 2.}), | |
| ('tensor', 'sum(W @ x)', {'W': np.array([[1., 2.], [3., 4.]]), 'x': np.array([2., 3.])}), | |
| ('tensor', 'sum(x @ W)', {'W': np.array([[1., 2., 3.], [3., 4., 5.]]), 'x': np.array([2., 3.])}), | |
| ('tensor', 'softmax(x)[-1]', {'x': np.array([1., 2., 3.])}), | |
| ('tensor', 'sum(x*x)', {'x': np.array([1., 2., 3.])}), | |
| ('tensor', 'sum(x*x)', {'x': 3.}), | |
| ] | |
| def test_every_view_uses_identical_local_equations(notation, engine, formula, feed): | |
| model = scalar(formula, feed) if engine == 'scalar' else tensor_graph(formula, feed).to_computation_model() | |
| text, html, diagram = render_trace(model, notation), render_trace_html(model, notation), render_diagram(model, notation) | |
| for step in model.steps: | |
| renderer = MathRenderer(notation, html=True) | |
| upper = renderer.backward_rule(step.explanation) | |
| lower = renderer.application_derivation(step.explanation, layout="auto") | |
| plain_upper, plain_lower = MathRenderer(notation).explanation(step.explanation) | |
| assert step_label(step, notation).split(DIVIDER)[0] == f"<div class='backward-rule'>{upper}</div>" | |
| assert f'>{lower}</div>' in step_label(step, notation) | |
| box = _contribution_box_html(step.operation, step, notation) | |
| assert f'class="tv-deriv-rule">{upper}</div>' in box | |
| assert f'class="tv-deriv-result">{lower}</div>' in box | |
| assert box in html | |
| assert upper in diagram and lower in diagram | |
| assert plain_upper in text and plain_lower in text | |
| assert 'diff(out,' not in upper and '∂out/' not in upper | |
| # Semantic identity: the lower compartment never re-states the symbolic rule. | |
| assert step.rule.application not in step.explanation.application_derivation.stages | |
| assert step.explanation.application_derivation.lhs == step.explanation.contribution | |
| # Expressions are data, independent of both display notation and containers. | |
| assert '<span' not in json.dumps(step.rule.application.to_data()) | |
| def test_reference_scalar_example_and_occurrence_identity(): | |
| model = scalar('(x*x)*y+(y+2)', {'x': 3., 'y': 4.}) | |
| x_steps = [s for s in model.steps if s.destination_quantity.name == 'x'] | |
| renderer = MathRenderer() | |
| assert renderer.explanation(x_steps[0].explanation) == ( | |
| 'a = x, b = x\n∂(a · b)/∂a = b = x = 3', | |
| 'c_n1→x, 1 = 4 · 3 = 12') | |
| assert renderer.explanation(x_steps[1].explanation) == ( | |
| 'a = x, b = x\n∂(a · b)/∂b = a = x = 3', | |
| 'c_n1→x, 2 = 4 · 3 = 12') | |
| assert x_steps[0].rule.operand.key != x_steps[1].rule.operand.key | |
| assert x_steps[0].destination_quantity is x_steps[1].destination_quantity | |
| assert '∂(x · x)/∂x' not in render_trace(model) | |
| assert 'Adjoint:\nx̄ = ∂out/∂x = 24' in render_trace(model) | |
| for step in x_steps: | |
| assert renderer.expression(step.explanation.contribution) in renderer.explanation(step.explanation)[1] | |
| assert MathRenderer(html=True).expression(step.explanation.contribution) in render_trace_html(model) | |
| def test_slot_names_do_not_confuse_graph_quantity_names(): | |
| model = scalar('a*a', {'a': 3.}) | |
| assert MathRenderer().explanation(model.steps[0].explanation)[0] == 'b = a, c = a\n∂(b · c)/∂b = c = a = 3' | |
| def test_contract_depends_on_local_shapes(): | |
| scalar_model = scalar('x*x', {'x': 3.}) | |
| tensor_scalar = tensor_graph('x*x', {'x': 3.}).to_computation_model() | |
| assert all(isinstance(s.rule, ScalarLocalRule) for s in tensor_scalar.steps) | |
| for first, second in zip(scalar_model.steps, tensor_scalar.steps): | |
| assert MathRenderer().explanation(first.explanation) == MathRenderer().explanation(second.explanation) | |
| dot = tensor_graph('dot(x,x)', {'x': np.array([2., 3.])}).to_computation_model() | |
| assert all(isinstance(s.rule, TensorBackwardRule) for s in dot.steps) | |
| assert all(s.contribution.shape == (2,) for s in dot.steps) | |
| def test_refactored_tensor_rules_against_finite_differences(formula, feed): | |
| graph = tensor_graph(formula, feed) | |
| for name, original in feed.items(): | |
| value = np.asarray(original) | |
| for index in np.ndindex(value.shape): | |
| plus = {k: np.array(v, copy=True) for k, v in feed.items()} | |
| minus = {k: np.array(v, copy=True) for k, v in feed.items()} | |
| plus[name][index] += 1e-6 | |
| minus[name][index] -= 1e-6 | |
| numerical = (tensor_graph(formula, plus, True).root.val - tensor_graph(formula, minus, True).root.val) / 2e-6 | |
| assert np.isclose(np.asarray(graph.name_to_input[name].grad)[index], numerical, atol=1e-7) | |
| model = graph.to_computation_model() | |
| values = {q.uid: q.primal_value for q in model.quantities} | |
| for step in model.steps: | |
| env = {**values, 'adjoint:' + step.source_quantity.uid: step.incoming_adjoint} | |
| actual = step.rule.application.evaluate(env) | |
| assert np.shape(actual) == step.destination_quantity.shape | |
| np.testing.assert_allclose(actual, step.contribution) | |
| def test_snapshot_is_independent_of_feed_and_reevaluation(): | |
| feed = {'x': np.array([2., 3.])} | |
| graph = tensor_graph('sum(x*x)', feed) | |
| model = graph.to_computation_model() | |
| before = render_trace(model) | |
| feed['x'][0] = 100. | |
| graph.evaluate({'x': np.array([10., 20.])}) | |
| assert render_trace(model) == before | |
| x = next(q for q in model.quantities if q.name == 'x') | |
| np.testing.assert_array_equal(x.primal_value, [2., 3.]) | |
| with pytest.raises(ValueError): | |
| x.primal_value[0] = 99. | |
| def test_rendering_does_not_evaluate_rules(monkeypatch): | |
| from math_expression import Expr | |
| model = tensor_graph('sum(W @ x)', {'W': np.ones((2, 3)), 'x': np.ones(3)}).to_computation_model() | |
| def forbidden(*args, **kwargs): | |
| raise AssertionError('Renderers must consume stored explanations') | |
| monkeypatch.setattr(Expr, 'evaluate', forbidden) | |
| render_diagram(model) | |
| render_trace(model) | |
| render_trace_html(model) | |
| def test_forward_only_html_does_not_construct_backward_examples(monkeypatch): | |
| from teaching_examples import multiplication_example | |
| multiplication_example.cache_clear() | |
| graph = ComputationGraph().build(safe_parse_expr('x*x'), forward_only=True) | |
| graph.evaluate({sp.Symbol('x'): 3.}, forward_only=True) | |
| def forbidden(*args, **kwargs): | |
| raise AssertionError('Forward-only rendering must not differentiate examples') | |
| monkeypatch.setattr(sp, 'diff', forbidden) | |
| html = render_trace_html(graph.to_computation_model(forward_only=True), mode='forward') | |
| body = html.split('</style>')[1] | |
| assert '∂' not in body and 'diff(' not in body | |
| assert 'tv-adjoint-step' not in body and 'tv-deriv-box' not in body | |
| def test_repeated_base_and_exponent_keep_different_rules(): | |
| model = scalar('x**x', {'x': 2.}) | |
| first, second = model.steps | |
| assert first.rule.operand.key != second.rule.operand.key | |
| assert first.local_value == 4. | |
| assert np.isclose(second.local_value, 4 * np.log(2)) | |
| assert 'log(' in MathRenderer('code').explanation(second.explanation)[0] | |
| assert np.isclose(model.accumulations[0].resulting_adjoint, 4 + 4 * np.log(2)) | |
| def test_numerical_outer_product_is_not_rendered_as_row_times_column(): | |
| model = tensor_graph('sum(W @ x)', {'W': np.ones((2, 2)), 'x': np.array([2., 3.])}).to_computation_model() | |
| upper, lower = MathRenderer().explanation(model.steps[0].explanation) | |
| assert upper == 'c_n1→W = n̄1 xᵀ' | |
| assert lower == 'c_n1→W = outer([1, 1], [2, 3])\n = [[2, 3], [2, 3]]' | |
| def test_constant_zero_power_does_not_multiply_infinite_adjoint(value): | |
| count = np.size(value) | |
| graph = tensor_graph(f'log(sum(x**0)-{count})', {'x': value}) | |
| np.testing.assert_array_equal(graph.name_to_input['x'].grad, np.zeros_like(value)) | |
| model = graph.to_computation_model() | |
| step = next(s for s in model.steps if s.operation.kind == 'pow') | |
| assert not step.rule.application.has_adjoint() | |
| assert 'nan' not in MathRenderer().explanation(step.explanation)[1] | |
| def test_unreached_adjoints_are_absent_but_active_zeros_remain(engine, formula, expected): | |
| from diagram_renderer import quantity_label | |
| from tikz_renderer import render_tikz | |
| feed = {'x': 0. if formula == 'x*x' else 3.} | |
| model = (scalar(formula, feed) if engine == 'scalar' else | |
| tensor_graph(formula, feed).to_computation_model()) | |
| x = next(q for q in model.quantities if q.name == 'x') | |
| assert x.adjoint_value == expected | |
| operations = {op.quantity.uid: op for op in model.operations} | |
| if 'nograd' in formula: | |
| stopped_operand = next(op.operands[0] for op in model.operations if op.function == 'nograd') | |
| assert stopped_operand.adjoint_value is None | |
| stop_result = next(op.quantity for op in model.operations if op.function == 'nograd') | |
| assert stop_result.adjoint_value == 1 | |
| assert 'adjoint-state' not in quantity_label(stopped_operand, operations[stopped_operand.uid], 'code') | |
| label = quantity_label(x, None, 'code') | |
| assert ('adjoint-state' in label) == (expected is not None) | |
| tikz = render_tikz(model) | |
| assert (r'\overline{x}=' in tikz) == (expected is not None) | |
| text = render_trace(model, 'code') | |
| assert ('diff(out, x) = x_bar =' in text) == (expected is not None) | |
| html = render_trace_html(model, 'code') | |
| assert 'None' not in html | |
| def test_semantic_derivation_stages_model_contract(): | |
| """Lower compartment = (concrete application, result); never the symbolic rule.""" | |
| def check(step, concrete_kind): | |
| d = step.explanation.application_derivation | |
| assert d.lhs == step.explanation.contribution | |
| assert len(d.stages) == 2 | |
| assert d.stages[0].kind == concrete_kind | |
| assert d.stages[1].kind == 'literal' | |
| # Identity, not substring: the symbolic rule lives only in local_rule. | |
| assert step.rule.application not in d.stages | |
| assert not set(step.explanation.local_rule.terms[1:]) & set(d.stages) | |
| return d.stages[1].value | |
| # 1. Scalar multiplication: c_out→x = 1 · 3 = 3 | |
| assert check(scalar('x*y', {'x': 2., 'y': 3.}).steps[0], 'mul') == 3.0 | |
| # 2. Repeated scalar operand: x*x | |
| for s in scalar('x*x', {'x': 3.}).steps: | |
| assert check(s, 'mul') == 3.0 | |
| # 3. softmax(z)[1]: softmax VJP, then index selection | |
| tg_sm = tensor_graph('softmax(z)[1]', {'z': np.array([1., 2., 3.])}).to_computation_model() | |
| np.testing.assert_allclose(check(tg_sm.steps[0], 'elementwise'), | |
| [-0.02203304, 0.18483645, -0.1628034], atol=1e-5) | |
| np.testing.assert_allclose(check(tg_sm.steps[1], 'scale'), [0., 1., 0.]) | |
| # 4. sum(W @ x): outer product to W, transpose-matmul to x | |
| tg_mv = tensor_graph('sum(W @ x)', {'W': np.ones((2, 3)), 'x': np.ones(3)}).to_computation_model() | |
| np.testing.assert_array_equal(check(tg_mv.steps[0], 'outer'), np.ones((2, 3))) | |
| np.testing.assert_array_equal(check(tg_mv.steps[1], 'matmul'), [2., 2., 2.]) | |
| # 5. Broadcasting: sum(x/y + b) | |
| tg_bc = tensor_graph('sum(x/y + b)', {'x': np.array([2., 4.]), 'y': np.array([1., 2.]), 'b': 2.}).to_computation_model() | |
| assert check(next(s for s in tg_bc.steps if s.destination_quantity.name == 'b'), 'sum') == 2.0 | |
| def test_derivation_stages_forbid_embedded_equality(): | |
| """Negative invariant: No derivation stage may itself be an Equation or encode internal equality.""" | |
| from backward_explanation import Equation | |
| for engine, formula, feed in CASES: | |
| model = scalar(formula, feed) if engine == 'scalar' else tensor_graph(formula, feed).to_computation_model() | |
| for step in model.steps: | |
| d = step.explanation.application_derivation | |
| # lhs must be an Expr, not Equation | |
| assert not isinstance(d.lhs, Equation) | |
| assert d.lhs.kind != 'equation' | |
| for stage in d.stages: | |
| # stages must be Expr, never Equation | |
| assert not isinstance(stage, Equation) | |
| assert stage.kind != 'equation' | |