"""Renderer-independent computation store and typed information flow.""" from dataclasses import dataclass, field from enum import Enum from typing import Any import numpy as np from math_expression import Expr, operation_expression, snapshot from backward_rules import ScalarLocalRule, TensorBackwardRule from backward_explanation import BackwardExplanation, explain_step, explain_stop from scalar_rules import scalar_rule_from_record class BackwardEdgeType(Enum): ADJOINT_INTRODUCTION = 'adjoint_introduction' ADJOINT_PROPAGATION = 'adjoint_propagation' CONTRIBUTION = 'contribution' @dataclass class QuantityState: uid: str name: str primal_value: Any adjoint_value: Any differentiable: bool is_input: bool shape: tuple is_constant: bool = False @dataclass class Operation: quantity: QuantityState operands: list[QuantityState] kind: str function: str | None = None exponent: float | None = None index: int | None = None formula: Expr = field(init=False) stop_explanation: BackwardExplanation | None = None def __post_init__(self): self.formula = operation_expression(self) @dataclass class BackwardStep: uid: str source_quantity: QuantityState destination_quantity: QuantityState incoming_adjoint: Any operation: Operation rule: ScalarLocalRule | TensorBackwardRule local_value: Any contribution: Any operand_occurrence: int explanation: BackwardExplanation | None = None @dataclass class Accumulation: uid: str destination_quantity: QuantityState contributions: list[BackwardStep] resulting_adjoint: Any @dataclass class BackwardEdge: source: str destination: str kind: BackwardEdgeType value: Any quantity: QuantityState | None = None # A single contribution is allowed to collapse into an introduction edge. contribution_step: BackwardStep | None = None seed: bool = False @dataclass class OutputState: uid: str quantity: QuantityState @dataclass class ForwardEdge: source: str destination: str value: Any operand_occurrence: int | None = None @dataclass class ComputationModel: quantities: list[QuantityState] operations: list[Operation] root: QuantityState output_name: str formula: str forward_only: bool output: OutputState | None = None steps: list[BackwardStep] = field(default_factory=list) accumulations: list[Accumulation] = field(default_factory=list) forward_edges: list[ForwardEdge] = field(default_factory=list) backward_edges: list[BackwardEdge] = field(default_factory=list) def build_computation_model(graph, engine, output_name='out', forward_only=False): """Adapt numerical engine records once; formatting does not infer edge roles.""" if engine not in {'scalar', 'tensor'}: raise ValueError('Unknown numerical engine') quantities, operations, steps = [], [], [] by_id = {} for node in graph.nodes: differentiable = not node.is_constant if node.is_input else any(by_id[a.uid].differentiable for a in node.args) q = QuantityState(node.uid, node.name, snapshot(node.val), None if forward_only or not differentiable or not node.live else snapshot(node.grad), differentiable, node.is_input, np.shape(node.val), is_constant=getattr(node, 'is_constant', False)) by_id[q.uid] = q quantities.append(q) if node.is_input: continue if engine == 'scalar': names = {'Add': 'add', 'Mul': 'mul', 'Pow': 'pow'} func = node.expr.func.__name__ operation = Operation(q, [by_id[a.uid] for a in node.args], names.get(func, 'func'), func) else: operation = Operation(q, [by_id[a.uid] for a in node.args], node.op_type, node.func_name, node.pow_exp, node.pick_index) operations.append(operation) is_nograd = func == 'nograd' if engine == 'scalar' else node.func_name == 'nograd' if is_nograd and not forward_only and q.differentiable: operation.stop_explanation = explain_stop(operation) if forward_only or not getattr(node, 'live', True) or is_nograd: continue records = node.local_derivatives if engine == 'scalar' else node.backward_rules for record in records: occurrence = record['operand_occurrence'] destination = by_id[record['arg_node'].uid] if not destination.differentiable: continue rule = (scalar_rule_from_record(operation, occurrence, record) if engine == 'scalar' else record['rule']) local_value = record['local_val'] if engine == 'scalar' else record.get('local_val') step = BackwardStep( f'der_{node.uid}_{occurrence}', q, destination, snapshot(node.grad), operation, rule, snapshot(local_value), snapshot(record['grad_contrib'] if engine == 'scalar' else record['contrib_val']), occurrence, ) values = {a.uid: a.primal_value for a in operation.operands} values[q.uid] = q.primal_value values['adjoint:' + q.uid] = step.incoming_adjoint step.explanation = explain_step(step, values) steps.append(step) root_q = by_id[graph.root.uid] output_state = OutputState(f'output_{root_q.uid}', root_q) model = ComputationModel(quantities, operations, root_q, output_name, str(graph.orig_expr) if engine == 'scalar' else graph.orig_expr_str, forward_only, output_state, steps) for operation in operations: for i, operand in enumerate(operation.operands): model.forward_edges.append(ForwardEdge(operand.uid, operation.quantity.uid, operand.primal_value, i)) # The final operation/result sends its primal result forward to the output presentation node model.forward_edges.append(ForwardEdge(model.root.uid, model.output.uid, model.root.primal_value, None)) if forward_only or not model.root.differentiable: return model model.backward_edges.append(BackwardEdge(model.output.uid, model.root.uid, BackwardEdgeType.ADJOINT_INTRODUCTION, model.root.adjoint_value, model.root, seed=True)) contributions = {q.uid: [] for q in quantities} for step in steps: contributions[step.destination_quantity.uid].append(step) model.backward_edges.append(BackwardEdge(step.source_quantity.uid, step.uid, BackwardEdgeType.ADJOINT_PROPAGATION, step.incoming_adjoint, step.source_quantity)) for quantity in quantities: incoming = contributions[quantity.uid] if len(incoming) > 1: accumulation = Accumulation('sum_' + quantity.uid, quantity, incoming, quantity.adjoint_value) model.accumulations.append(accumulation) for step in incoming: model.backward_edges.append(BackwardEdge(step.uid, accumulation.uid, BackwardEdgeType.CONTRIBUTION, step.contribution, contribution_step=step)) model.backward_edges.append(BackwardEdge(accumulation.uid, quantity.uid, BackwardEdgeType.ADJOINT_INTRODUCTION, quantity.adjoint_value, quantity)) elif incoming: step = incoming[0] model.backward_edges.append(BackwardEdge(step.uid, quantity.uid, BackwardEdgeType.ADJOINT_INTRODUCTION, quantity.adjoint_value, quantity, step)) return model def tensor_operation(node): """Numerical tensor-node adapter; rule factories depend only on this record.""" def quantity(n): return QuantityState(n.uid, n.name, n.val, n.grad, not n.is_constant, n.is_input, n.shape, n.is_constant) return Operation(quantity(node), [quantity(a) for a in node.args], node.op_type, node.func_name, node.pow_exp, node.pick_index)