Spaces:
Running
Running
File size: 3,312 Bytes
b9499ab | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 | import type { Graph, GraphNode, Mode, Operation } from "./types";
// Rounded for reading
export function formatNumber(value: number): string {
return String(Number(value.toPrecision(5)));
}
// A multi-part argument, such as the constant "2 * 3", in brackets
function atom(arg: string): string {
return arg.includes(" ") ? `(${arg})` : arg;
}
// Each operation's derivative with respect to each argument; backend/src/logic.py's PARTIALS as text
const PARTIALS: Record<Operation, ((a: string, b: string) => string)[]> = {
add: [() => "1", () => "1"],
sub: [() => "1", () => "−1"],
mul: [(_, b) => b, (a) => a],
div: [(_, b) => `1 / ${b}`, (a, b) => `−${a} / ${b}^2`],
pow: [(a, b) => `${b} · ${a}^(${b} − 1)`, (a, b) => `${a}^${b} · log(${a})`],
neg: [() => "−1"],
sin: [(a) => `cos(${a})`],
cos: [(a) => `−sin(${a})`],
tan: [(a) => `1 / cos(${a})^2`],
exp: [(a) => `exp(${a})`],
log: [(a) => `1 / ${a}`],
sqrt: [(a) => `1 / (2 · sqrt(${a}))`],
};
export function partialOf(op: Operation, index: number, args: string[]): string {
return PARTIALS[op][index](atom(args[0]), atom(args[1] ?? ""));
}
// An id with a bar over its first letter: the adjoint of that node
export function adjoint(id: string): string {
return `${id[0]}̄${id.slice(1)}`;
}
// An id with a dot over its first letter: the tangent of that node
function tangent(id: string): string {
return `${id[0]}\u0307${id.slice(1)}`;
}
export function derivativeSymbol(mode: Mode, id: string): string {
return mode === "forward" ? tangent(id) : adjoint(id);
}
// A derivative times a local partial, leaving out a factor of 1
function term(symbol: string, partial: string): string {
if (partial === "1") {
return symbol;
}
if (partial === "−1") {
return `−${symbol}`;
}
return `${symbol} · ${partial.includes(" ") ? `(${partial})` : partial}`;
}
function sumOf(terms: string[]): string {
return terms.join(" + ").replaceAll("+ −", "− ");
}
// How reverse mode gets a node's adjoint: one term for each use of the node
export function adjointFormula(graph: Graph, node: GraphNode): string {
if (node.id === graph.output) {
return `${adjoint(node.id)} = 1`;
}
const terms: string[] = [];
for (const consumer of graph.nodes) {
if (consumer.kind === "op") {
consumer.args.forEach((arg, i) => {
if (arg === node.id) {
terms.push(term(adjoint(consumer.id), partialOf(consumer.op, i, consumer.args)));
}
});
}
}
return `${adjoint(node.id)} = ${sumOf(terms)}`;
}
// A node's arguments that have derivatives, in the order of a step's partials; constants don't
export function nodeArgs(graph: Graph, node: GraphNode & { kind: "op" }): { arg: string; index: number }[] {
return node.args.map((arg, index) => ({ arg, index })).filter(({ arg }) => !(arg in graph.constants));
}
// How forward mode gets a node's tangent: one term for each of its arguments
export function tangentFormula(graph: Graph, node: GraphNode, seed: string): string {
if (node.kind === "input") {
return `${tangent(node.id)} = ${node.name === seed ? 1 : 0}`;
}
const terms = nodeArgs(graph, node).map(({ arg, index }) => term(tangent(arg), partialOf(node.op, index, node.args)));
return `${tangent(node.id)} = ${sumOf(terms)}`;
}
|