"""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 # ============================================================================== @pytest.mark.parametrize("expr,vals,factory", [ ("x * y", {"x": 2.0, "y": 3.0}, build_scalar_model), ("(x * x) * y + (y + 2)", {"x": 3.0, "y": 4.0}, build_scalar_model), ("sum(W @ x)", {"W": np.array([[1.0, 2.0], [3.0, 4.0]]), "x": np.array([1.0, 2.0])}, build_tensor_model), ("softmax(W @ x + b)[1]", { "W": np.array([[0.5, -0.2], [-0.1, 0.8], [0.3, 0.1]]), "x": np.array([1.0, 2.0]), "b": np.array([0.1, -0.1, 0.0]) }, build_tensor_model), ]) 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 '
' in html assert '
' 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 '
' in html assert '
' in html # 3. Individual step cards must NOT be open by default (Headlines view) assert '
' not in html assert '
' 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('Backward pass') 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 '
' 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