'
''
)
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"
Stop backward propagation through {escape(operand.name)}.
')
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('seed')
if is_nograd:
tags.append('stop-gradient boundary')
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'
Primal
'
f'{_primal_header_html(q, op, notation, format_config=format_config)}']
if len(received) > 1:
body_parts.append('')
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"
from {escape(producer.name)}{suffix}: {c_sub} = {_esc_val(s.contribution, format_config=format_config)}
")
accumulate = " + ".join(_esc_val(s.contribution, format_config=format_config) for s in received)
body_parts.append('
')
_TV_NOTATION_LEGEND = """
How to read the notation
{{primal}}
Primal — the value computed in the forward pass. Read from left to right: quantity = computation = values used = result.
{{adjoint}}
Adjoint — sensitivity of the final output.n1 tells us how sensitive the final output is to n1, including all downstream paths.
{{rule}}
{{application}}
Contribution — one local part of an adjoint. The first line gives the local derivative. The second combines it with the incoming adjoint to give this operand’s contribution to x. Here a and b name independent operand slots, both filled by x.
{{accumulation}}
When there is more than one contribution. Contributions from every use of a quantity are added to establish its adjoint.
Colors follow meaning: green = primal, violet = adjoint, magenta = contribution, rose = local derivative, gray = constant.
"""
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 ('How to read the notation'
f'
')
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
"""
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 = (
'
'
''
''
''
''
'
'
)
if forward_only:
toolbar = toolbar.replace("Forward / Backward only", "Forward only")
parts = [_TV_STYLE, '
']
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'{escape(q.name)} = {_esc_val(q.primal_value, format_config=format_config)}'
for q in inputs)
forward_cards = "".join(_forward_card_html(op, notation, format_config=format_config) for op in model.operations)
parts.append(
'Forward pass'
'
'
f'
Inputs
{input_chips}
'
f'{forward_cards}
'
)
if forward_only or not model.root.differentiable:
parts.append('
')
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 = '
Click ⊙ focus on an adjoint below to trace only the quantities that contribute to it; click it again to clear.