Spaces:
Running
Running
| """Typed local rules shared by numerical evaluation and all presentations.""" | |
| from dataclasses import dataclass, replace | |
| import numpy as np | |
| from math_expression import Expr, adjoint, literal, node, operation_expression, ref | |
| class ScalarLocalRule: | |
| formula: Expr | |
| operand: Expr | |
| derivative: Expr | |
| mapped_derivative: Expr | |
| application: Expr | |
| bindings: tuple[tuple[Expr, Expr], ...] = () | |
| notes: tuple[str, ...] = () | |
| class TensorBackwardRule: | |
| formula: Expr | |
| operand: Expr | |
| application: Expr | |
| notes: tuple[str, ...] = () | |
| def scalar_rule(operation, occurrence, derivative, slot_derivative=None, slots=()): | |
| repeated = sum(q.uid == operation.operands[occurrence].uid for q in operation.operands) > 1 | |
| args = tuple(ref(q) for q in operation.operands) | |
| bindings = tuple(zip(slots, args)) if repeated else () | |
| return ScalarLocalRule( | |
| operation_expression(operation, slots if repeated else args), | |
| slots[occurrence] if repeated else args[occurrence], | |
| slot_derivative if repeated else derivative, | |
| derivative, node('mul', adjoint(operation.quantity), derivative), bindings) | |
| def tensor_rule(operation, occurrence): | |
| """Construct one local pullback. Its expression is also its implementation.""" | |
| from scalar_rules import operand_slots | |
| slots = operand_slots(operation) | |
| args = slots | |
| shapes = [q.shape for q in operation.operands] | |
| child = operation.operands[occurrence] | |
| g, y = adjoint(operation.quantity), ref(operation.quantity) | |
| x = args[occurrence] | |
| kind, function = operation.kind, operation.function | |
| notes = () | |
| factor = None | |
| product = 'elementwise' if operation.quantity.shape else 'scale' | |
| if kind == 'matmul': | |
| a, b = args | |
| ranks = tuple(map(len, shapes)) | |
| if ranks == (1, 1): | |
| application = node('scale', g, args[1 - occurrence]) | |
| elif ranks == (2, 1): | |
| application = node('outer', g, b) if occurrence == 0 else node('matmul', node('transpose', a), g) | |
| elif ranks == (1, 2): | |
| application = node('matmul', b, g) if occurrence == 0 else node('outer', a, g) | |
| elif ranks == (2, 2): | |
| application = node('matmul', g, node('transpose', b)) if occurrence == 0 else node('matmul', node('transpose', a), g) | |
| else: | |
| raise ValueError(f'Unsupported matrix product ranks: {ranks}') | |
| elif function == 'softmax': | |
| application = node('elementwise', y, node('sub', g, node('sum', node('elementwise', g, y)))) | |
| else: | |
| if kind == 'neg': | |
| factor = literal(-1) | |
| elif kind == 'add': | |
| factor = literal(1) | |
| elif kind == 'sub': | |
| factor = literal(1 if occurrence == 0 else -1) | |
| elif kind in {'mul', 'dot'}: | |
| factor = args[1 - occurrence] | |
| elif kind == 'div': | |
| factor = node('div', literal(1), args[1]) if occurrence == 0 else node('div', node('neg', args[0]), node('pow', args[1], literal(2))) | |
| elif kind == 'pow': | |
| p = operation.exponent | |
| factor = (node('zeros_like', x) if child.shape else literal(0)) if p == 0 else node('mul', literal(p), node('pow', x, literal(p - 1))) | |
| elif function == 'pick': | |
| factor = node('one_hot', x, literal(operation.index)) | |
| notes = ('Indices are zero-based; negative indices count from the end.',) | |
| elif function in {'sum', 'mean'}: | |
| factor = node('ones_like', x) if child.shape else literal(1) | |
| if function == 'mean' and child.shape: | |
| factor = node('div', factor, node('size', x)) | |
| notes = (f'N = number of entries = {int(np.prod(child.shape))}',) | |
| elif function in {'max', 'min'}: | |
| mask = node('equal', x, y) | |
| factor = node('div', mask, node('sum', mask)) if child.shape else literal(1) | |
| notes = ('Equal-share subgradient at ties.',) | |
| elif function == 'sigmoid': | |
| factor = node('elementwise' if child.shape else 'mul', y, node('sub', literal(1), y)) | |
| elif function == 'relu': | |
| factor = node('positive', x) | |
| notes = ('At zero, the chosen derivative is 0.',) | |
| elif function == 'tanh': | |
| factor = node('sub', literal(1), node('pow', y, literal(2))) | |
| elif function == 'exp': | |
| factor = y | |
| elif function == 'log': | |
| factor = node('div', literal(1), x) | |
| if factor is None: | |
| raise ValueError(f'Missing backward rule for {kind}/{function}') | |
| # A zero exponent is constant even if the incoming adjoint is nonfinite. | |
| application = factor if kind == 'pow' and operation.exponent == 0 else node(product, g, factor) | |
| if not child.shape and operation.quantity.shape: | |
| application = node('sum', application) | |
| notes += ('Sum over broadcast entries.',) | |
| bindings = {slot.key: ref(q) for slot, q in zip(slots, operation.operands)} | |
| application = application.substitute(bindings) | |
| if not operation.quantity.shape and not any(shapes) and factor is not None: | |
| # Scalar operations inside a tensor graph use the same scalar contract. | |
| rule = scalar_rule(operation, occurrence, factor.substitute(bindings), factor, slots) | |
| return replace(rule, application=application, notes=notes) | |
| return TensorBackwardRule(operation.formula, ref(child), application, notes) | |