Spaces:
Sleeping
Sleeping
| """Unit and snapshot tests for the stepwise text trace renderer.""" | |
| import re | |
| import numpy as np | |
| import pytest | |
| import sympy as sp | |
| from autodiff_visualizer import ComputationGraph, safe_parse_expr | |
| from tensor_engine import TensorComputationGraph | |
| from trace_renderer import render_trace, render_trace_html | |
| def build_scalar_model(expr_str, values, forward_only=False): | |
| g = ComputationGraph().build(safe_parse_expr(expr_str), forward_only=forward_only) | |
| g.evaluate({sp.Symbol(k): float(v) for k, v in values.items()}, forward_only=forward_only) | |
| return g.to_diagram_model(forward_only=forward_only) | |
| def build_tensor_model(expr_str, values, forward_only=False): | |
| g = TensorComputationGraph().build_from_ast(expr_str) | |
| g.evaluate(values, forward_only=forward_only) | |
| return g.to_diagram_model(forward_only=forward_only) | |
| # ============================================================================== | |
| # 1. Forward-only mode purity tests | |
| # ============================================================================== | |
| def test_forward_only_mode_purity(expr, vals, factory): | |
| model = factory(expr, vals, forward_only=True) | |
| trace = render_trace(model, notation="mathbook", mode="forward") | |
| # Invariants for Forward-only: | |
| # 1. Header is FORWARD PASS | |
| assert "FORWARD PASS" in trace | |
| assert "Inputs" in trace | |
| assert "OUTPUT" in trace | |
| # 2. Absolutely NO backward pass tokens | |
| assert "BACKWARD PASS" not in trace | |
| assert "Seed" not in trace | |
| assert "Backpropagate" not in trace | |
| assert "contribution" not in trace | |
| assert "ACCUMULATION" not in trace | |
| assert "FINAL GRADIENTS" not in trace | |
| assert "∂" not in trace | |
| assert "diff(" not in trace | |
| assert "̄" not in trace # No overbar combining character | |
| assert "_bar" not in trace | |
| assert "—" not in trace # No em-dash placeholder | |
| # ============================================================================== | |
| # 2. Semantic tests: Step structure, accumulation timing, and backward pass | |
| # ============================================================================== | |
| def test_scalar_canonical_x_times_y(): | |
| model = build_scalar_model("x * y", {"x": 2.0, "y": 3.0}) | |
| trace = render_trace(model, notation="mathbook", mode="forward_backward") | |
| assert "FORWARD PASS" in trace | |
| assert "BACKWARD PASS" in trace | |
| assert "Seed" in trace | |
| assert "out̄ = ∂out/∂out = 1" in trace | |
| # Output operation out = x · y | |
| assert "out = x · y" in trace | |
| assert "out = 6" in trace | |
| # Backprop through out: | |
| assert "Backpropagate through out" in trace | |
| assert "To x:" in trace | |
| assert "∂(x · y)/∂x = y = 3" in trace | |
| assert "c_out→x = 1 · 3 = 3" in trace | |
| assert "out̄ · y" not in trace | |
| assert "Contribution to x:\n3" in trace | |
| assert "Adjoint:\nx̄ = ∂out/∂x = 3" in trace | |
| assert "To y:" in trace | |
| assert "∂(x · y)/∂y = x = 2" in trace | |
| assert "c_out→y = 1 · 2 = 2" in trace | |
| assert "Contribution to y:\n2" in trace | |
| assert "Adjoint:\nȳ = ∂out/∂y = 2" in trace | |
| # Final gradients | |
| assert "FINAL GRADIENTS" in trace | |
| assert "∂out/∂x = x̄ = 3" in trace | |
| assert "∂out/∂y = ȳ = 2" in trace | |
| def test_scalar_x_times_x_accumulation(): | |
| """For x*x, there are 2 operand occurrences of x. Accumulation must add both contributions.""" | |
| model = build_scalar_model("x * x", {"x": 3.0}) | |
| trace = render_trace(model, notation="mathbook", mode="forward_backward") | |
| # Both operand arguments represented | |
| assert "Argument 1: x" in trace | |
| assert "Contribution to x:\n3" in trace | |
| # Accumulation occurs | |
| assert "x̄ = 3 + 3 = 6" in trace or "x̄ = 6" in trace | |
| assert "Adjoint:\nx̄ = ∂out/∂x = 6" in trace | |
| def test_scalar_shared_intermediate_accumulation(): | |
| """For (x+y)*(x+y), accumulation of n1 must occur before backpropagating through n1.""" | |
| model = build_scalar_model("(x + y) * (x + y)", {"x": 2.0, "y": 1.0}) | |
| trace = render_trace(model, notation="mathbook", mode="forward_backward") | |
| # n1 is shared intermediate | |
| assert "n1 = x + y" in trace | |
| # Accumulation of n1 | |
| accum_pos = trace.find("n̄1 =") | |
| backprop_n1_pos = trace.find("Backpropagate through n1") | |
| assert accum_pos != -1 | |
| assert backprop_n1_pos != -1 | |
| # Crucial semantic invariant: accumulation occurs BEFORE backpropagation through n1 | |
| assert accum_pos < backprop_n1_pos | |
| def test_scalar_fixture_expression(): | |
| """Canonical expression (x*x)*y + (y+2) with x=3, y=4.""" | |
| model = build_scalar_model("(x * x) * y + (y + 2)", {"x": 3.0, "y": 4.0}) | |
| trace = render_trace(model, notation="mathbook", mode="forward_backward") | |
| # Result should be 42 | |
| assert "out = 42" in trace | |
| # Adjoints should be x̄ = 24, ȳ = 10 | |
| assert "∂out/∂x = x̄ = 24" in trace | |
| assert "∂out/∂y = ȳ = 10" in trace | |
| # In-place accumulation for y preserves arrival order (1 then 9) | |
| assert "ȳ = 1 + 9" in trace and "= 10" in trace | |
| assert "x̄ = 12 + 12" in trace and "= 24" in trace | |
| assert "Adjoint:\nȳ = ∂out/∂y = 10" in trace | |
| assert "Adjoint:\nx̄ = ∂out/∂x = 24" in trace | |
| assert "Therefore:" not in trace | |
| assert "contribution =" not in trace | |
| # ============================================================================== | |
| # 3. Notation tests: Mathbook vs Code | |
| # ============================================================================== | |
| def test_mathbook_vs_code_notation(): | |
| model = build_scalar_model("x * y", {"x": 2.0, "y": 3.0}) | |
| trace_math = render_trace(model, notation="mathbook", mode="forward_backward") | |
| trace_code = render_trace(model, notation="code", mode="forward_backward") | |
| # Mathbook uses ∂ and overbar | |
| assert "∂" in trace_math | |
| assert "out̄" in trace_math | |
| assert "x̄" in trace_math | |
| # Code uses diff and _bar | |
| assert "diff(" in trace_code | |
| assert "out_bar" in trace_code | |
| assert "x_bar" in trace_code | |
| assert "∂" not in trace_code | |
| assert "̄" not in trace_code | |
| # ============================================================================== | |
| # 4. Tensor operations and budget invariant | |
| # ============================================================================== | |
| def test_tensor_trace_matrix_vector(): | |
| W = np.array([[1.0, 2.0], [3.0, 4.0]]) | |
| x = np.array([1.0, 2.0]) | |
| model = build_tensor_model("sum(W @ x)", {"W": W, "x": x}) | |
| trace = render_trace(model, notation="mathbook", mode="forward_backward") | |
| # Value invariant: No "shape(" substituted for values | |
| assert "shape(" not in trace | |
| # Tensor values are formatted cleanly | |
| assert "[5, 11]" in trace or "[5.0, 11.0]" in trace or "5" in trace | |
| assert "out = 16" in trace | |
| # Backward pass present | |
| assert "BACKWARD PASS" in trace | |
| assert "out̄ = 1" in trace | |
| def test_tensor_trace_softmax(): | |
| z = np.array([1.0, 2.0, 3.0]) | |
| model = build_tensor_model("softmax(z)[1]", {"z": z}) | |
| trace = render_trace(model, notation="mathbook", mode="forward_backward") | |
| # Values must be present as numbers | |
| assert "shape(" not in trace | |
| assert "FORWARD PASS" in trace | |
| assert "BACKWARD PASS" in trace | |
| assert "Seed" in trace | |
| assert "FINAL GRADIENTS" in trace | |
| # ============================================================================== | |
| # 5. Contract tests: Architectural decoupling & formatting invariants | |
| # ============================================================================== | |
| def test_no_diagram_renderer_import_in_trace_renderer(): | |
| """Trace renderer must be decoupled from diagram renderer.""" | |
| import sys | |
| import trace_renderer | |
| with open(trace_renderer.__file__, "r", encoding="utf-8") as f: | |
| content = f.read() | |
| assert "diagram_renderer" not in content, "trace_renderer.py must not import diagram_renderer" | |
| assert "_clean_math_text" not in content, "trace_renderer.py must not rely on HTML-stripping regexes" | |
| def test_constants_distinct_from_inputs(): | |
| """Constants must not appear in the Inputs section and must not render as '2 = 2'.""" | |
| model = build_scalar_model("(x * x) * y + (y + 2)", {"x": 3.0, "y": 4.0}) | |
| trace = render_trace(model, notation="mathbook", mode="forward_backward") | |
| # Inputs section only has named inputs x and y | |
| inputs_section = trace.split("1. ")[0] | |
| assert "x = 3" in inputs_section | |
| assert "y = 4" in inputs_section | |
| assert "2 = 2" not in inputs_section | |
| assert "cst" not in inputs_section | |
| def test_forward_computed_primal_equation(): | |
| """Forward step header must introduce computed equation once: res = formula = value.""" | |
| model = build_scalar_model("x * y", {"x": 2.0, "y": 3.0}) | |
| trace = render_trace(model, notation="mathbook", mode="forward") | |
| assert "1. out = x · y = 6" in trace | |
| assert "Apply:" in trace | |
| assert "2 · 3" in trace | |
| assert "= 6" in trace | |
| # No redundant "Result:" section | |
| assert "Result:" not in trace | |
| def test_matmul_trace_vjp_rules(): | |
| """Matmul VJP rules must be rendered properly in both mathbook and code notation.""" | |
| W = np.array([[0.5, -0.2], [-0.1, 0.8]]) | |
| x = np.array([1.0, 2.0]) | |
| model = build_tensor_model("sum(W @ x)", {"W": W, "x": x}) | |
| trace_math = render_trace(model, notation="mathbook", mode="forward_backward") | |
| assert "Backpropagate through n1 = W @ x" in trace_math | |
| assert "xᵀ" in trace_math or "Wᵀ" in trace_math | |
| assert "Local backward rule:" in trace_math | |
| assert "Apply:" in trace_math | |
| assert "Contribution to W:" in trace_math | |
| assert "Contribution to x:" in trace_math | |
| trace_code = render_trace(model, notation="code", mode="forward_backward") | |
| assert "diff(" in trace_code | |
| assert "@" in trace_code or ".T" in trace_code | |
| def test_html_trace_formula_and_equations_card(): | |
| """HTML trace must render the selected formula and n-notation equations before the Forward pass box.""" | |
| model = build_scalar_model("(x * x) * y + (y + 2)", {"x": 3.0, "y": 4.0}) | |
| html = render_trace_html(model, notation="mathbook") | |
| assert '<details class="tv-formula-card">' in html | |
| assert '<details class="tv-formula-card" open>' not in html | |
| assert 'Selected Formula' in html | |
| assert 'Decomposed Equations' in html | |
| # Formula card must appear BEFORE Forward pass details | |
| formula_idx = html.find('class="tv-formula-card"') | |
| forward_idx = html.find('class="tv-region tv-forward-region"') | |
| assert formula_idx != -1 and forward_idx != -1 | |
| assert formula_idx < forward_idx | |
| # Equations list contains n-notation | |
| assert 'n1 = x · x = 9' in html | |
| assert 'n2 = n1 · y = 36' in html | |
| assert 'n3 = y + 2 = 6' in html | |
| assert 'out = n2 + n3 = 42' in html | |
| def test_html_trace_tensor_formula_and_equations(): | |
| """HTML trace for tensor operations must also show formula and n-notation decomposition.""" | |
| W = np.array([[0.5, -0.2], [-0.1, 0.8]]) | |
| x = np.array([1.0, 2.0]) | |
| model = build_tensor_model("sum(W @ x)", {"W": W, "x": x}) | |
| html = render_trace_html(model, notation="mathbook") | |
| assert 'class="tv-formula-card"' in html | |
| assert 'out = sum(W @ x)' in html or 'sum(W @ x)' in html | |
| assert 'n1 = W @ x' in html | |
| assert 'out = ∑(n1)' in html | |
| def test_html_trace_defaults_to_headlines_view(): | |
| """HTML trace must default to Level 2 (Headlines) with step cards collapsed and regions open.""" | |
| model = build_scalar_model("(x * x) * y + (y + 2)", {"x": 3.0, "y": 4.0}) | |
| html = render_trace_html(model, notation="mathbook") | |
| # 1. Toolbar active button must be Headlines (level 2) | |
| assert 'class="tv-level-btn tv-active" data-level="2"' in html | |
| assert 'class="tv-level-btn tv-active" data-level="3"' not in html | |
| # 2. Regions must be open | |
| assert '<details class="tv-region tv-forward-region" open>' in html | |
| assert '<details class="tv-region tv-backward-region" open>' in html | |
| # 3. Individual step cards must NOT be open by default (Headlines view) | |
| assert '<details class="tv-step tv-primal-step"' in html | |
| assert '<details class="tv-step tv-primal-step" data-q="n1" open>' not in html | |
| assert '<details class="tv-step tv-adjoint-step" data-q="n1" open>' not in html | |
| # 4. Focus hint must appear inside the Backward pass region body as the first line | |
| bwd_region_idx = html.rfind('tv-backward-region') | |
| bwd_summary_idx = html.rfind('<summary>Backward pass</summary>') | |
| hint_idx = html.rfind('tv-focus-hint') | |
| assert bwd_region_idx != -1 and bwd_summary_idx != -1 and hint_idx != -1 | |
| assert bwd_summary_idx < hint_idx | |
| def test_html_trace_notation_legend_learner_centric(): | |
| """HTML trace must omit top color legend and have conceptual learner-centric notation help.""" | |
| model = build_scalar_model("(x * x) * y + (y + 2)", {"x": 3.0, "y": 4.0}) | |
| html = render_trace_html(model, notation="mathbook") | |
| # 1. Top color legend is removed from normal trace view | |
| assert 'class="tv-legend"' not in html | |
| # 2. Notation help is present and learner-centric | |
| assert '<details class="tv-notation">' in html | |
| assert 'How to read the notation' in html | |
| assert 'Primal — the value computed in the forward pass' in html | |
| assert 'Read from left to right: quantity = computation = values used = result' in html | |
| assert 'Adjoint — sensitivity of the final output' in html | |
| assert 'tells us how sensitive the final output is' in html | |
| assert 'Contribution — one local part of an adjoint' in html | |
| assert 'When there is more than one contribution' in html | |
| assert 'Contributions from every use of a quantity are added to establish its adjoint' in html | |
| assert 'Colors follow meaning: green = primal, violet = adjoint, magenta = contribution, rose = local derivative, gray = constant' in html | |
| # 3. Implementation jargon must be absent | |
| assert 'CSS' not in html | |
| assert 'card header' not in html | |
| assert 'heading disambiguates' not in html | |
| assert 'chip' not in html.lower() or 'input_chips' in html # no "Constant chip" | |
| assert 'Constant chip' not in html | |
| def test_trace_with_custom_format_config(): | |
| """Verify trace rendering respects format_config for precision and scientific notation.""" | |
| from number_format import NumberFormatConfig | |
| W = np.array([[0.51234, -0.21234], [-0.11234, 0.81234]]) | |
| x = np.array([1.0001, 2.0002]) | |
| model = build_tensor_model("sum(W @ x)", {"W": W, "x": x}) | |
| # Render with 4 decimals | |
| cfg4 = NumberFormatConfig(precision=4, allow_scientific=False) | |
| trace_text = render_trace(model, notation="mathbook", format_config=cfg4) | |
| assert "0.5123" in trace_text | |
| assert "1.0001" in trace_text | |
| # Render HTML with 1 decimal | |
| cfg1 = NumberFormatConfig(precision=1, allow_scientific=False) | |
| trace_html = render_trace_html(model, notation="mathbook", format_config=cfg1) | |
| assert "[0.5, -0.2]" in trace_html | |
| assert "[1, 2]" in trace_html | |