Spaces:
Running
Running
Simon Clematide
Format walkthrough forward operands with MathRenderer and remove Operands label
769c8c2 | """Stepwise text trace renderer consuming ComputationModel. | |
| The trace linearizes forward evaluation and reverse-mode autodiff, | |
| producing a structured plaintext trace in execution order. | |
| """ | |
| from html import escape | |
| import json | |
| import re | |
| from math_expression import Expr, adjoint, literal, node, ref | |
| from math_renderer import MathRenderer | |
| from backward_rules import ScalarLocalRule | |
| from backward_explanation import Equation, contribution_symbol, adjoint_equation | |
| from teaching_examples import multiplication_example | |
| from tensor_formatter import format_tensor_value, DEFAULT_TENSOR_ELEMENT_BUDGET | |
| from number_format import format_numeric_literals | |
| def _format_val(val, config=None): | |
| if val is None: | |
| return "—" | |
| return format_tensor_value(val, element_budget=DEFAULT_TENSOR_ELEMENT_BUDGET, include_shape_if_truncated=True, config=config) | |
| def _adj_name(q, notation, format_config=None): | |
| return MathRenderer(notation, format_config=format_config).expression(adjoint(q)) | |
| def _partial_sym(source_name, target_name, notation, format_config=None): | |
| return MathRenderer(notation, format_config=format_config).expression(node('global_derivative', | |
| Expr('ref', key='output', name=source_name), Expr('ref', key='target', name=target_name))) | |
| def _op_formula_str(op, notation, format_config=None): | |
| return MathRenderer(notation, format_config=format_config).expression(op.formula) | |
| def _op_applied_str(op, notation, format_config=None): | |
| values = {q.uid: q.primal_value for q in op.operands} | |
| return MathRenderer(notation, format_config=format_config).expression(op.formula.bind(values)) | |
| def render_trace(model, notation='mathbook', mode='forward_backward', format_config=None): | |
| """Linearizes model into deterministic sequential execution steps. | |
| notation: 'mathbook' | 'code' | |
| mode: 'forward' | 'forward_backward' | |
| """ | |
| if notation not in {'mathbook', 'code'}: | |
| raise ValueError('Unknown notation') | |
| forward_only = (mode == 'forward' or model.forward_only) | |
| lines = [] | |
| # ========================================================================= | |
| # FORWARD PASS | |
| # ========================================================================= | |
| lines.append("FORWARD PASS") | |
| lines.append("=" * 12) | |
| lines.append("") | |
| lines.append("Inputs") | |
| lines.append("------") | |
| inputs = [q for q in model.quantities if q.is_input and not q.is_constant] | |
| for inp in inputs: | |
| lines.append(f"{inp.name} = {_format_val(inp.primal_value, config=format_config)}") | |
| lines.append("") | |
| step_idx = 1 | |
| for op in model.operations: | |
| res = op.quantity | |
| header_text = f"{step_idx}. {res.name} = {_op_formula_str(op, notation, format_config=format_config)} = {_format_val(res.primal_value, config=format_config)}" | |
| lines.append(header_text) | |
| lines.append("-" * len(header_text)) | |
| lines.append("Operands:") | |
| for operand in op.operands: | |
| if operand.is_constant: | |
| lines.append(f"{_format_val(operand.primal_value, config=format_config)}") | |
| else: | |
| lines.append(f"{operand.name} = {_format_val(operand.primal_value, config=format_config)}") | |
| lines.append("Apply:") | |
| lines.append(f"{_op_applied_str(op, notation, format_config=format_config)}") | |
| lines.append(f"= {_format_val(res.primal_value, config=format_config)}") | |
| lines.append("") | |
| step_idx += 1 | |
| lines.append("OUTPUT") | |
| lines.append("------") | |
| lines.append(f"{model.output_name} = {_format_val(model.root.primal_value, config=format_config)}") | |
| lines.append("") | |
| if forward_only or not model.root.differentiable: | |
| return "\n".join(lines).strip() | |
| # ========================================================================= | |
| # BACKWARD PASS | |
| # ========================================================================= | |
| lines.append("BACKWARD PASS") | |
| lines.append("=" * 13) | |
| lines.append("") | |
| lines.append("Seed") | |
| lines.append("----") | |
| root_adj = _adj_name(model.root, notation, format_config=format_config) | |
| root_deriv = _partial_sym(model.output_name, model.root.name, notation, format_config=format_config) | |
| lines.append(f"{root_adj} = {root_deriv} = {_format_val(model.root.adjoint_value, config=format_config)}") | |
| lines.append("") | |
| steps_by_source = {} | |
| for s in model.steps: | |
| steps_by_source.setdefault(s.source_quantity.uid, []).append(s) | |
| acc_by_dest = {acc.destination_quantity.uid: acc for acc in model.accumulations} | |
| contributions_received = {q.uid: [] for q in model.quantities} | |
| # Reverse topological execution through operations | |
| bwd_idx = step_idx | |
| for op in reversed(model.operations): | |
| res_q = op.quantity | |
| if res_q.uid not in steps_by_source: | |
| continue | |
| op_steps = steps_by_source[res_q.uid] | |
| if not op_steps: | |
| continue | |
| bwd_header = f"{bwd_idx}. Backpropagate through {res_q.name} = {_op_formula_str(op, notation, format_config=format_config)}" | |
| lines.append(bwd_header) | |
| lines.append("-" * len(bwd_header)) | |
| lines.append("Incoming adjoint:") | |
| lines.append(f"{_adj_name(res_q, notation, format_config=format_config)} = {_format_val(op_steps[0].incoming_adjoint, config=format_config)}") | |
| lines.append("") | |
| for step in op_steps: | |
| target = step.destination_quantity | |
| occ = step.operand_occurrence | |
| occ_header = f"Argument {occ + 1}: {target.name}" if len(op.operands) > 1 and len(set(q.uid for q in op.operands)) == 1 else f"To {target.name}:" | |
| lines.append(occ_header) | |
| upper, lower = MathRenderer(notation, format_config=format_config).explanation(step.explanation) | |
| lines.append("Local derivative:" if isinstance(step.rule, ScalarLocalRule) else "Local backward rule:") | |
| lines.append(upper) | |
| lines.append("Apply:") | |
| lines.append(lower) | |
| for note in step.explanation.notes: | |
| lines.append(f"Note: {note}") | |
| lines.append(f"Contribution to {target.name}:") | |
| lines.append(f"{_format_val(step.contribution, config=format_config)}") | |
| lines.append("") | |
| # Record contribution | |
| contributions_received[target.uid].append(step) | |
| # Check if all contributions for this target have arrived and accumulation can be emitted | |
| if target.uid in acc_by_dest: | |
| acc = acc_by_dest[target.uid] | |
| if len(contributions_received[target.uid]) == len(acc.contributions): | |
| acc_header = f"Accumulate contributions to {target.name}" | |
| lines.append(acc_header) | |
| lines.append("-" * len(acc_header)) | |
| # Use actual order of arrival: contributions_received[target.uid] | |
| contrib_vals = [_format_val(s.contribution, config=format_config) for s in contributions_received[target.uid]] | |
| contrib_sum_expr = " + ".join(contrib_vals) | |
| lines.append(f"{_adj_name(target, notation, format_config=format_config)} = {contrib_sum_expr}") | |
| lines.append(f"= {_format_val(acc.resulting_adjoint, config=format_config)}") | |
| lines.append("Adjoint:") | |
| root_name = model.output_name | |
| deriv_name = _partial_sym(root_name, target.name, notation, format_config=format_config) | |
| lines.append(f"{_adj_name(target, notation, format_config=format_config)} = {deriv_name} = {_format_val(acc.resulting_adjoint, config=format_config)}") | |
| lines.append("") | |
| elif len(contributions_received[target.uid]) == 1: | |
| # Single contribution establishing adjoint | |
| all_steps_for_target = [s for s in model.steps if s.destination_quantity.uid == target.uid] | |
| if len(all_steps_for_target) == 1: | |
| lines.append("Adjoint:") | |
| root_name = model.output_name | |
| deriv_name = _partial_sym(root_name, target.name, notation, format_config=format_config) | |
| lines.append(f"{_adj_name(target, notation, format_config=format_config)} = {deriv_name} = {_format_val(step.contribution, config=format_config)}") | |
| lines.append("") | |
| bwd_idx += 1 | |
| # Final gradients summary for input quantities | |
| diff_inputs = [q for q in model.quantities if q.is_input and q.adjoint_value is not None] | |
| if diff_inputs: | |
| lines.append("FINAL GRADIENTS") | |
| lines.append("=" * 15) | |
| root_name = model.output_name | |
| for inp in diff_inputs: | |
| partial = _partial_sym(root_name, inp.name, notation, format_config=format_config) | |
| adj = _adj_name(inp, notation, format_config=format_config) | |
| lines.append(f"{partial} = {adj} = {_format_val(inp.adjoint_value, config=format_config)}") | |
| return "\n".join(lines).strip() | |
| # ============================================================================= | |
| # HTML trace view | |
| # | |
| # Exactly two top-level box types, reusing the diagram's semantic color | |
| # contract: a Primal box (green, rounded) summarizes an established forward | |
| # value as one equation (quantity = symbolic rule = substituted rule = value); | |
| # an Adjoint box (violet, rectangular) summarizes an established total output | |
| # sensitivity and contains the complete local derivation needed to explain it. | |
| # Contributions, local derivative rules, and accumulation are explicit inside | |
| # Adjoint boxes but never become peer top-level boxes. | |
| # | |
| # Ownership: an operation's card PRODUCES contributions toward its operands | |
| # (shown as "Contribution to TARGET" subsections); a quantity's own card | |
| # RECEIVES contributions, owns accumulation, and stores the adjoint (the card | |
| # header). One quantity, one Adjoint box, regardless of how many operations | |
| # consume it. | |
| # ============================================================================= | |
| _TV_STYLE = """ | |
| <style> | |
| .tv-root { font-family: ui-monospace, "SFMono-Regular", Menlo, Consolas, monospace; color: #202124; } | |
| .tv-root * { box-sizing: border-box; } | |
| .tv-region { border: 1px solid #d8dbe0; border-radius: 10px; margin-bottom: 16px; overflow: hidden; } | |
| .tv-region > summary { | |
| list-style: none; cursor: pointer; padding: 10px 14px; font-weight: 650; | |
| font-size: 0.92rem; font-family: sans-serif; background: #fbfbfa; border-bottom: 1px solid #d8dbe0; | |
| } | |
| .tv-region > summary::-webkit-details-marker { display: none; } | |
| .tv-region.tv-forward-region > summary { color: #16803a; } | |
| .tv-region.tv-backward-region > summary { color: #7a3e9d; } | |
| .tv-region-body { padding: 14px; } | |
| .tv-step-label { color: #5b5f66; font-size: 0.72rem; font-family: sans-serif; margin-bottom: 2px; } | |
| .tv-inputs-list, .tv-operand-list { display: flex; flex-direction: column; align-items: flex-start; gap: 2px; } | |
| .tv-inputs-list { margin: 2px 0 14px 0; } | |
| .tv-step { | |
| border: 1px solid #d8dbe0; margin-bottom: 10px; background: #fff; | |
| } | |
| .tv-step:last-child { margin-bottom: 0; } | |
| .tv-step > summary { | |
| list-style: none; cursor: pointer; padding: 8px 12px; display: flex; width: 100%; | |
| align-items: baseline; gap: 8px; font-size: 0.88rem; | |
| } | |
| .tv-step > summary::-webkit-details-marker { display: none; } | |
| .tv-step-body { padding: 6px 12px 12px 20px; display: flex; flex-direction: column; align-items: flex-start; gap: 10px; } | |
| .tv-primal-step { border-color: #16803a; background: #f2f8f3; border-radius: 12px; width: fit-content; max-width: 100%; } | |
| .tv-primal-step > summary { color: #16803a; font-weight: 650; } | |
| .tv-adjoint-step { border-color: #7a3e9d; background: #fff; border-radius: 3px; } | |
| .tv-adjoint-step > summary { color: #7a3e9d; font-weight: 650; } | |
| .tv-primal-box { | |
| display: inline-block; width: fit-content; max-width: 100%; | |
| background: #f2f8f3; border: 1px solid #16803a; border-radius: 6px; padding: 5px 10px; font-size: 0.9rem; | |
| } | |
| .tv-deriv-box { border: 1px solid #b83e72; border-radius: 6px; overflow-x: auto; font-size: 0.85rem; width: fit-content; max-width: 100%; } | |
| .tv-deriv-rule { background: #fbeaf1; color: #202124; padding: 4px 10px; } | |
| .tv-deriv-result { background: #f9ebf2; color: #932f6d; padding: 4px 10px; border-top: 1px solid #b83e72; } | |
| .tv-forward { color: #16803a; } | |
| .tv-adjoint { color: #7a3e9d; } | |
| .tv-contribution { color: #932f6d; } | |
| .tv-constant { display: inline-block; background: #f2f2f2; border: 1px solid #777; border-radius: 4px; padding: 1px 7px; } | |
| .tv-constant-plain { color: #777; } | |
| .tv-ov { text-decoration: overline; text-decoration-thickness: 1.4px; } | |
| .tv-hr { border: none; border-top: 1px solid #d8dbe0; margin: 2px 0; width: 100%; } | |
| .tv-tag { font-size: 0.78rem; color: #5b5f66; font-style: italic; font-family: sans-serif; } | |
| .tv-effect { margin: 2px 0 0 0; font-size: 0.85rem; color: #5b5f66; font-family: sans-serif; } | |
| .tv-toolbar { display: flex; flex-wrap: wrap; align-items: center; gap: 8px; margin-bottom: 8px; font-family: sans-serif; } | |
| .tv-toolbar button { | |
| all: unset !important; box-sizing: border-box !important; cursor: pointer !important; | |
| font-family: sans-serif !important; font-size: 0.78rem !important; font-weight: 400 !important; | |
| line-height: 1.2 !important; padding: 5px 11px !important; border-radius: 6px !important; | |
| border: 1px solid #d8dbe0 !important; background: #fbfbfa !important; color: #202124 !important; | |
| display: inline-block !important; | |
| } | |
| .tv-toolbar button:hover { border-color: #5b5f66 !important; } | |
| .tv-toolbar button.tv-active { background: #202124 !important; color: #fff !important; border-color: #202124 !important; } | |
| .tv-toolbar .tv-fullscreen-btn { margin-left: auto !important; } | |
| .tv-focus-hint { font-size: 0.76rem; color: #5b5f66; font-family: sans-serif; margin: 0 0 14px; } | |
| .tv-focus-btn { | |
| all: unset !important; box-sizing: border-box !important; cursor: pointer !important; | |
| margin-left: auto !important; flex: none !important; display: inline-block !important; | |
| font-family: sans-serif !important; font-size: 0.72rem !important; font-weight: 400 !important; | |
| line-height: 1.2 !important; padding: 3px 8px !important; border-radius: 5px !important; | |
| border: 1px solid transparent !important; background: transparent !important; color: #5b5f66 !important; | |
| } | |
| .tv-focus-btn:hover { border-color: #7a3e9d !important; color: #7a3e9d !important; } | |
| [data-q] { transition: opacity 0.15s ease, filter 0.15s ease, box-shadow 0.15s ease; } | |
| .tv-dimmed { opacity: 0.25; filter: grayscale(70%); } | |
| .tv-step.tv-trace-highlighted { box-shadow: 0 0 0 2px #7a3e9d; } | |
| .tv-notation { border: 1px solid #d8dbe0; border-radius: 8px; background: #fbfbfa; margin-bottom: 16px; overflow: hidden; font-family: sans-serif; } | |
| .tv-notation > summary { | |
| list-style: none; cursor: pointer; padding: 10px 14px; font-size: 0.82rem; font-weight: 650; | |
| } | |
| .tv-notation > summary::-webkit-details-marker { display: none; } | |
| .tv-notation-grid { padding: 4px 14px 14px; display: flex; flex-direction: column; gap: 12px; border-top: 1px solid #d8dbe0; } | |
| .tv-notation-row { display: grid; grid-template-columns: minmax(180px, 280px) 1fr; gap: 16px; align-items: center; font-size: 0.82rem; } | |
| .tv-notation-row p { margin: 0; color: #5b5f66; line-height: 1.45; font-family: sans-serif; } | |
| .tv-notation-row p code { color: #202124; font-family: ui-monospace, monospace; } | |
| .tv-notation-ex { font-size: 0.84rem; font-family: ui-monospace, monospace; } | |
| .tv-rule-note { font-size:0.85em; color:#5b5f66; margin:4px 0; } | |
| .tv-notation-color-note { margin: 4px 0 0 0; padding-top: 10px; border-top: 1px solid #e2e5e9; font-size: 0.78rem; color: #5b5f66; font-family: sans-serif; font-style: italic; } | |
| @media (max-width: 600px) { .tv-notation-row { grid-template-columns: 1fr; gap: 4px; } } | |
| .tv-formula-card { | |
| border: 1px solid #d8dbe0; border-radius: 10px; background: #fbfbfa; | |
| margin-bottom: 16px; font-family: sans-serif; overflow: hidden; | |
| } | |
| .tv-formula-card > summary { | |
| list-style: none; cursor: pointer; padding: 10px 14px; font-weight: 650; | |
| font-size: 0.92rem; font-family: sans-serif; background: #fbfbfa; color: #202124; | |
| display: flex; align-items: center; justify-content: space-between; | |
| } | |
| .tv-formula-card > summary::-webkit-details-marker { display: none; } | |
| .tv-formula-card[open] > summary { border-bottom: 1px solid #d8dbe0; } | |
| .tv-formula-card-body { padding: 12px 14px; } | |
| .tv-formula-target { | |
| font-family: ui-monospace, "SFMono-Regular", Menlo, Consolas, monospace; | |
| font-size: 1.05rem; font-weight: 650; color: #202124; margin: 4px 0 10px; | |
| word-break: break-word; | |
| } | |
| .tv-equations-list { | |
| display: flex; flex-direction: column; gap: 4px; padding: 4px 0 2px; | |
| } | |
| .tv-equation-item { | |
| font-family: ui-monospace, "SFMono-Regular", Menlo, Consolas, monospace; | |
| font-size: 0.88rem; color: #16803a; word-break: break-word; | |
| } | |
| .tv-root:fullscreen { background: #fff; padding: 16px; overflow: auto; } | |
| </style> | |
| """ | |
| def _esc_val(v, format_config=None): | |
| return escape(_format_val(v, config=format_config)) | |
| def _adj_symbol_html(q, notation, format_config=None): | |
| return MathRenderer(notation, html=True, format_config=format_config).expression(adjoint(q)) | |
| def _format_formula_str(formula_str, notation, format_config=None): | |
| if not formula_str: | |
| return "" | |
| s = format_numeric_literals(str(formula_str), config=format_config) | |
| if notation == 'mathbook': | |
| s = s.replace('**', '^') | |
| s = s.replace('*', ' · ') | |
| return s | |
| def _formula_card_html(model, notation, format_config=None): | |
| """Renders the selected formula and intermediate decomposed equations (n-notation).""" | |
| formula_disp = _format_formula_str(model.formula, notation, format_config=format_config) | |
| target_line = f"{escape(model.output_name)} = {escape(formula_disp)}" if formula_disp else escape(model.output_name) | |
| eq_lines = [] | |
| for op in model.operations: | |
| q_name = escape(op.quantity.name) | |
| formula = MathRenderer(notation, html=True, format_config=format_config).expression(op.formula) | |
| val = _esc_val(op.quantity.primal_value, format_config=format_config) | |
| eq_lines.append(f'<div class="tv-equation-item">{q_name} = {formula} = {val}</div>') | |
| summary_text = f"Selected Formula: {target_line}" | |
| return ( | |
| '<details class="tv-formula-card">' | |
| f'<summary><span>{summary_text}</span></summary>' | |
| '<div class="tv-formula-card-body">' | |
| '<div class="tv-step-label">Selected Formula</div>' | |
| f'<div class="tv-formula-target">{target_line}</div>' | |
| '<div class="tv-step-label" style="margin-top:10px;">Decomposed Equations</div>' | |
| f'<div class="tv-equations-list">{"".join(eq_lines)}</div>' | |
| '</div>' | |
| '</details>' | |
| ) | |
| def _primal_header_html(q, op, notation, format_config=None): | |
| """quantity = value, or quantity = symbolic rule = substituted rule = value.""" | |
| if op is None: | |
| return f"{escape(q.name)} = {_esc_val(q.primal_value, format_config=format_config)}" | |
| renderer = MathRenderer(notation, html=True, format_config=format_config) | |
| applied = op.formula.bind({a.uid: a.primal_value for a in op.operands}) | |
| return (f"{escape(q.name)} = {renderer.expression(op.formula)}" | |
| f" = {renderer.expression(applied)} = {_esc_val(q.primal_value, format_config=format_config)}") | |
| def _repeated_operand_suffix(op, step): | |
| """' (operand k)' only when the target occurs more than once among this operation's operands.""" | |
| occurrences = [i for i, a in enumerate(op.operands) if a.uid == step.destination_quantity.uid] | |
| if len(occurrences) > 1: | |
| return f" (operand {step.operand_occurrence + 1})" | |
| return "" | |
| def _forward_card_html(op, notation, format_config=None): | |
| res = op.quantity | |
| header = _primal_header_html(res, op, notation, format_config=format_config) | |
| renderer = MathRenderer(notation, html=True, format_config=format_config) | |
| operand_lines = [] | |
| for i, operand in enumerate(op.operands): | |
| if operand.is_constant: | |
| val_html = renderer.expression(literal(operand.primal_value)) | |
| operand_lines.append( | |
| f"<div class='tv-constant-plain'>Operand {i + 1}: {val_html}</div>") | |
| else: | |
| eq_html = f"{renderer.expression(ref(operand))} = {renderer.expression(literal(operand.primal_value))}" | |
| operand_lines.append( | |
| f"<div class='tv-forward'>Operand {i + 1}: {eq_html}</div>") | |
| return (f'<details class="tv-step tv-primal-step" data-q="{escape(res.name)}"><summary><span>{header}</span></summary>' | |
| f'<div class="tv-step-body">' | |
| f'<div class="tv-operand-list">{"".join(operand_lines)}</div></div></details>') | |
| def _contribution_box_html(op, step, notation, format_config=None): | |
| """Views wrap shared mathematical equations; they never rebuild local rules.""" | |
| renderer = MathRenderer(notation, html=True, format_config=format_config) | |
| rule_line = renderer.backward_rule(step.explanation) | |
| result_line = renderer.application_derivation(step.explanation) | |
| notes = ''.join(f'<div class="tv-rule-note">{escape(note)}</div>' for note in step.explanation.notes) | |
| return (f'<div class="tv-deriv-box"><div class="tv-deriv-rule">{rule_line}</div>' | |
| f'<div class="tv-deriv-result">{result_line}</div></div>{notes}') | |
| def _nograd_effect_html(op, q, notation, format_config=None): | |
| operand = op.operands[0] | |
| rule_line, result_line = MathRenderer(notation, html=True, format_config=format_config).explanation(op.stop_explanation) | |
| return (f'<div><div class="tv-step-label">Contribution to {escape(operand.name)}</div>' | |
| f'<div class="tv-deriv-box"><div class="tv-deriv-rule">{rule_line}</div>' | |
| f'<div class="tv-deriv-result">{result_line}</div></div>' | |
| f'<div class="tv-step-label" style="margin-top:6px;">Effect</div>' | |
| f'<p class="tv-effect">Stop backward propagation through {escape(operand.name)}.</p></div>') | |
| def _backward_card_html(q, op, notation, model, contributions_received, steps_by_source, format_config=None): | |
| is_root = q.uid == model.root.uid | |
| is_nograd = op is not None and op.function == 'nograd' | |
| received = contributions_received.get(q.uid, []) | |
| tags = [] | |
| if is_root: | |
| tags.append('<span class="tv-tag">seed</span>') | |
| if is_nograd: | |
| tags.append('<span class="tv-tag">stop-gradient boundary</span>') | |
| tags_html = "".join(tags) | |
| header = MathRenderer(notation, html=True, format_config=format_config).equation(adjoint_equation(q, model.output_name, q.adjoint_value)) | |
| body_parts = [f'<div class="tv-step-label">Primal</div>' | |
| f'<span class="tv-primal-box">{_primal_header_html(q, op, notation, format_config=format_config)}</span>'] | |
| if len(received) > 1: | |
| body_parts.append('<div class="tv-hr"></div>') | |
| rec_lines = [] | |
| op_by_uid = {o.quantity.uid: o for o in model.operations} | |
| for s in received: | |
| producer = s.source_quantity | |
| producer_op = op_by_uid.get(producer.uid) | |
| suffix = _repeated_operand_suffix(producer_op, s) if producer_op else "" | |
| c_sub = MathRenderer(notation, html=True, format_config=format_config).expression(contribution_symbol(s)) | |
| rec_lines.append(f"<div class='tv-contribution'>from {escape(producer.name)}{suffix}: {c_sub} = {_esc_val(s.contribution, format_config=format_config)}</div>") | |
| accumulate = " + ".join(_esc_val(s.contribution, format_config=format_config) for s in received) | |
| body_parts.append('<div><div class="tv-step-label">Contributions received</div>' | |
| + "".join(rec_lines) + | |
| '<div class="tv-step-label" style="margin-top:6px;">Accumulate</div>' | |
| f'<span class="tv-contribution">{_adj_symbol_html(q, notation, format_config=format_config)} = {accumulate} = {_esc_val(q.adjoint_value, format_config=format_config)}</span></div>') | |
| producing_steps = steps_by_source.get(q.uid, []) | |
| if is_nograd: | |
| body_parts.append('<div class="tv-hr"></div>') | |
| body_parts.append(_nograd_effect_html(op, q, notation, format_config=format_config)) | |
| elif producing_steps: | |
| body_parts.append('<div class="tv-hr"></div>') | |
| for i, step in enumerate(producing_steps): | |
| target = step.destination_quantity | |
| suffix = _repeated_operand_suffix(op, step) | |
| margin = ' style="margin-top:8px;"' if i > 0 else '' | |
| body_parts.append(f'<div{margin}><div class="tv-step-label">Contribution to {escape(target.name)}{suffix}</div>' | |
| f'{_contribution_box_html(op, step, notation, format_config=format_config)}</div>') | |
| focus_btn = (f'<button type="button" class="tv-focus-btn" ' | |
| f'title="Show only what contributes to this adjoint" ' | |
| f'onclick="event.preventDefault();event.stopPropagation();tvFocusAdjoint(this,\'{escape(q.name)}\')">⊙ focus</button>') | |
| summary = f'<summary><span>{header}</span>{tags_html}{focus_btn}</summary>' | |
| return (f'<details class="tv-step tv-adjoint-step" data-q="{escape(q.name)}">{summary}' | |
| f'<div class="tv-step-body">{"".join(body_parts)}</div></details>') | |
| _TV_NOTATION_LEGEND = """ | |
| <details class="tv-notation"> | |
| <summary>How to read the notation</summary> | |
| <div class="tv-notation-grid"> | |
| <div class="tv-notation-row"> | |
| <span class="tv-primal-box tv-notation-ex">{{primal}}</span> | |
| <p><strong>Primal — the value computed in the forward pass.</strong> Read from left to right: quantity = computation = values used = result.</p> | |
| </div> | |
| <div class="tv-notation-row"> | |
| <span class="tv-notation-ex" style="border:1px solid #7a3e9d;border-radius:3px;padding:5px 10px;display:inline-block;width:fit-content;color:#7a3e9d;background:#fff;">{{adjoint}}</span> | |
| <p><strong>Adjoint — sensitivity of the final output.</strong> <span class="tv-ov">n1</span> tells us how sensitive the final output is to <i>n1</i>, including all downstream paths.</p> | |
| </div> | |
| <div class="tv-notation-row"> | |
| <div class="tv-deriv-box tv-notation-ex"> | |
| <div class="tv-deriv-rule">{{rule}}</div> | |
| <div class="tv-deriv-result">{{application}}</div> | |
| </div> | |
| <p><strong>Contribution — one local part of an adjoint.</strong> The first line gives the local derivative. The second combines it with the incoming adjoint to give this operand’s contribution to <i>x</i>. Here a and b name independent operand slots, both filled by x.</p> | |
| </div> | |
| <div class="tv-notation-row"> | |
| <div class="tv-notation-ex tv-contribution" style="padding:4px 0;">{{accumulation}}</div> | |
| <p><strong>When there is more than one contribution.</strong> Contributions from every use of a quantity are added to establish its adjoint.</p> | |
| </div> | |
| <div class="tv-notation-color-note"> | |
| Colors follow meaning: green = primal, violet = adjoint, magenta = contribution, rose = local derivative, gray = constant. | |
| </div> | |
| </div> | |
| </details> | |
| """ | |
| def _notation_legend(notation, forward_only=False): | |
| model = multiplication_example(forward_only) | |
| op = model.operations[0] | |
| renderer = MathRenderer(notation, html=True) | |
| primal = Equation((ref(op.quantity), op.formula, | |
| op.formula.bind({q.uid: q.primal_value for q in op.operands}), | |
| literal(op.quantity.primal_value))) | |
| if forward_only: | |
| return ('<details class="tv-notation"><summary>How to read the notation</summary>' | |
| f'<div class="tv-primal-box">{renderer.equation(primal)}</div>' | |
| '<p>Quantity = operation = substituted values = result.</p></details>') | |
| upper, lower = renderer.explanation(model.steps[0].explanation) | |
| acc = model.accumulations[0] | |
| accumulation = Equation((adjoint(acc.destination_quantity), | |
| node('add', *(contribution_symbol(s) for s in acc.contributions)), | |
| node('add', *(literal(s.contribution) for s in acc.contributions)), | |
| literal(acc.resulting_adjoint))) | |
| values = {'primal': renderer.equation(primal), | |
| 'adjoint': renderer.equation(adjoint_equation(op.quantity, model.output_name, op.quantity.adjoint_value)), | |
| 'rule': upper, 'application': lower, 'accumulation': renderer.equation(accumulation)} | |
| result = _TV_NOTATION_LEGEND | |
| for key, value in values.items(): | |
| result = result.replace('{{' + key + '}}', value) | |
| return result | |
| # gr.HTML sets innerHTML on update, and browsers never execute a <script> tag | |
| # inserted that way. This must be passed once via gr.HTML(..., head=TRACE_VIEW_SCRIPT) | |
| # at component-creation time instead of being embedded in the per-render string. | |
| TRACE_VIEW_SCRIPT = """ | |
| <script> | |
| (function() { | |
| window.tvSetLevel = function(level, btn, scope) { | |
| scope.querySelectorAll('details.tv-region').forEach(function(d) { d.open = level >= 2; }); | |
| scope.querySelectorAll('details.tv-step').forEach(function(d) { d.open = level >= 3; }); | |
| scope.querySelectorAll('.tv-level-btn').forEach(function(b) { b.classList.remove('tv-active'); }); | |
| btn.classList.add('tv-active'); | |
| }; | |
| window.tvFocusAdjoint = function(btn, q) { | |
| var scope = btn.closest('.tv-root'); | |
| var producers = JSON.parse(scope.getAttribute('data-tv-producers') || '{}'); | |
| var current = scope.getAttribute('data-tv-focused') || ''; | |
| var els = scope.querySelectorAll('[data-q]'); | |
| if (current === q) { | |
| scope.removeAttribute('data-tv-focused'); | |
| els.forEach(function(el) { el.classList.remove('tv-dimmed', 'tv-trace-highlighted'); }); | |
| return; | |
| } | |
| scope.setAttribute('data-tv-focused', q); | |
| var keep = new Set([q]); | |
| var stack = [q]; | |
| while (stack.length) { | |
| var cur = stack.pop(); | |
| (producers[cur] || []).forEach(function(p) { if (!keep.has(p)) { keep.add(p); stack.push(p); } }); | |
| } | |
| var headlinesBtn = scope.querySelector('.tv-level-btn[data-level="2"]'); | |
| if (headlinesBtn) tvSetLevel(2, headlinesBtn, scope); | |
| els.forEach(function(el) { | |
| var isKept = keep.has(el.getAttribute('data-q')); | |
| el.classList.toggle('tv-dimmed', !isKept); | |
| el.classList.toggle('tv-trace-highlighted', isKept); | |
| }); | |
| }; | |
| window.tvToggleFullscreen = function(btn) { | |
| var scope = btn.closest('.tv-root'); | |
| if (!document.fullscreenElement) { | |
| (scope.requestFullscreen || scope.webkitRequestFullscreen || function(){}).call(scope); | |
| } else { | |
| (document.exitFullscreen || document.webkitExitFullscreen || function(){}).call(document); | |
| } | |
| }; | |
| })(); | |
| </script> | |
| """ | |
| def render_trace_html(model, notation='mathbook', mode='forward_backward', format_config=None): | |
| """Renders the same ComputationModel as an HTML trace: exactly two box types | |
| (Primal, Adjoint), reusing the diagram's semantic color contract.""" | |
| if notation not in {'mathbook', 'code'}: | |
| raise ValueError('Unknown notation') | |
| forward_only = (mode == 'forward' or model.forward_only) | |
| toolbar = ( | |
| '<div class="tv-toolbar">' | |
| '<button type="button" class="tv-level-btn" data-level="1" onclick="tvSetLevel(1, this, this.closest(\'.tv-root\'))">Forward / Backward only</button>' | |
| '<button type="button" class="tv-level-btn tv-active" data-level="2" onclick="tvSetLevel(2, this, this.closest(\'.tv-root\'))">Headlines</button>' | |
| '<button type="button" class="tv-level-btn" data-level="3" onclick="tvSetLevel(3, this, this.closest(\'.tv-root\'))">Everything</button>' | |
| '<button type="button" class="tv-fullscreen-btn" onclick="tvToggleFullscreen(this)" title="Toggle fullscreen">⛶ Fullscreen</button>' | |
| '</div>' | |
| ) | |
| if forward_only: | |
| toolbar = toolbar.replace("Forward / Backward only", "Forward only") | |
| parts = [_TV_STYLE, '<div class="tv-root">'] | |
| parts.append(_notation_legend(notation, forward_only or not model.root.differentiable)) | |
| parts.append(toolbar) | |
| parts.append(_formula_card_html(model, notation, format_config=format_config)) | |
| # Forward pass | |
| inputs = [q for q in model.quantities if q.is_input and not q.is_constant] | |
| input_chips = "".join( | |
| f'<span class="tv-primal-box" data-q="{escape(q.name)}">{escape(q.name)} = {_esc_val(q.primal_value, format_config=format_config)}</span>' | |
| for q in inputs) | |
| forward_cards = "".join(_forward_card_html(op, notation, format_config=format_config) for op in model.operations) | |
| parts.append( | |
| '<details class="tv-region tv-forward-region" open><summary>Forward pass</summary>' | |
| '<div class="tv-region-body">' | |
| f'<div class="tv-step-label">Inputs</div><div class="tv-inputs-list">{input_chips}</div>' | |
| f'{forward_cards}</div></details>' | |
| ) | |
| if forward_only or not model.root.differentiable: | |
| parts.append('</div>') | |
| return "".join(parts) | |
| # Backward pass: quantity-indexed cards, in reverse construction order | |
| # (a valid reverse-topological order, since build_computation_model appends | |
| # each quantity only after every quantity it depends on). | |
| steps_by_source = {} | |
| for s in model.steps: | |
| steps_by_source.setdefault(s.source_quantity.uid, []).append(s) | |
| contributions_received = {} | |
| for s in model.steps: | |
| contributions_received.setdefault(s.destination_quantity.uid, []).append(s) | |
| op_by_uid = {op.quantity.uid: op for op in model.operations} | |
| seen, cards = set(), [] | |
| producers = {} | |
| for q in reversed(model.quantities): | |
| is_root = q.uid == model.root.uid | |
| if q.uid in seen or not (is_root or q.uid in contributions_received): | |
| continue | |
| seen.add(q.uid) | |
| cards.append(_backward_card_html(q, op_by_uid.get(q.uid), notation, model, contributions_received, steps_by_source, format_config=format_config)) | |
| producers[q.name] = sorted({s.source_quantity.name for s in contributions_received.get(q.uid, [])}) | |
| focus_hint = '<p class="tv-focus-hint">Click <strong>⊙ focus</strong> on an adjoint below to trace only the quantities that contribute to it; click it again to clear.</p>' | |
| parts.append( | |
| '<details class="tv-region tv-backward-region" open><summary>Backward pass</summary>' | |
| f'<div class="tv-region-body">{focus_hint}{"".join(cards)}</div></details>' | |
| ) | |
| parts.append('</div>') | |
| html = "".join(parts) | |
| producers_json = escape(json.dumps(producers), quote=True) | |
| return html.replace('<div class="tv-root">', f'<div class="tv-root" data-tv-producers="{producers_json}">', 1) | |