Spaces:
Running
Running
| """Parse a typed function into a computational graph, and run autodiff on it. | |
| A graph is {"inputs": [...], "constants": {...}, "nodes": [...], "output": id}. | |
| Nodes come in evaluation order, numbered w1, w2, ... with the inputs first. A | |
| node's id is what labels the edges leaving it: | |
| {"id": "w1", "kind": "input", "name": "x"} | |
| {"id": "w2", "kind": "op", "op": "mul", "args": ["2", "w1"]} | |
| An argument is a node's id or a key of constants, e.g. {"2": 2.0}. As in PyTorch | |
| and JAX, constants are part of an operation rather than nodes of the graph, and | |
| an expression of constants alone is worked out first: 2*3*x has one operation. | |
| Identical subexpressions share one node. Python's ast is used rather than | |
| sympy because sympy rewrites what was typed (x*x becomes x**2). | |
| A run is a list of steps, one per iteration of the stepper: | |
| {"phase": "forward" | "backward", "node": id, "value": v, "derivative": d, | |
| "partials": [{"arg": id, "partial": p, "amount": a}]} | |
| partial is the node's local derivative with respect to arg. In forward mode, | |
| derivative is the node's tangent, the sum of the amounts (partial * arg's tangent). | |
| In reverse mode's backward phase, derivative is the node's adjoint, and each | |
| amount (adjoint * partial) is added to arg's adjoint. Steps in its forward phase | |
| have no "derivative" key: Pyodide would turn a nested None into undefined, not null. | |
| """ | |
| import ast | |
| import math | |
| import re | |
| FUNCTIONS = ("sin", "cos", "tan", "exp", "log", "sqrt") | |
| CONSTANTS = {"pi": math.pi, "e": math.e} | |
| BINARY_OPS = {ast.Add: "add", ast.Sub: "sub", ast.Mult: "mul", ast.Div: "div", ast.Pow: "pow"} | |
| VALUES = { | |
| "add": lambda a, b: a + b, | |
| "sub": lambda a, b: a - b, | |
| "mul": lambda a, b: a * b, | |
| "div": lambda a, b: a / b, | |
| "pow": math.pow, | |
| "neg": lambda a: -a, | |
| "sin": math.sin, | |
| "cos": math.cos, | |
| "tan": math.tan, | |
| "exp": math.exp, | |
| "log": math.log, | |
| "sqrt": math.sqrt, | |
| } | |
| def _pow_by_base(a: float, b: float) -> float: | |
| return 0.0 if b == 0 else b * math.pow(a, b - 1) | |
| def _pow_by_exponent(a: float, b: float) -> float: | |
| # The limit at a = 0, as PyTorch uses | |
| return 0.0 if a == 0 and b >= 0 else math.pow(a, b) * math.log(a) | |
| # Each operation's derivative with respect to each argument | |
| PARTIALS = { | |
| "add": (lambda a, b: 1.0, lambda a, b: 1.0), | |
| "sub": (lambda a, b: 1.0, lambda a, b: -1.0), | |
| "mul": (lambda a, b: b, lambda a, b: a), | |
| "div": (lambda a, b: 1 / b, lambda a, b: -a / b**2), | |
| "pow": (_pow_by_base, _pow_by_exponent), | |
| "neg": (lambda a: -1.0,), | |
| "sin": (math.cos,), | |
| "cos": (lambda a: -math.sin(a),), | |
| "tan": (lambda a: 1 / math.cos(a) ** 2,), | |
| "exp": (math.exp,), | |
| "log": (lambda a: 1 / a,), | |
| "sqrt": (lambda a: 1 / (2 * math.sqrt(a)),), | |
| } | |
| _INFIX = {"add": "{:g} + {:g}", "sub": "{:g} - {:g}", "mul": "{:g} * {:g}", "div": "{:g} / {:g}", "pow": "{:g} ^ {:g}", "neg": "-{:g}"} | |
| def _formula(op: str, args: list[float]) -> str: | |
| return _INFIX.get(op, op + "({:g})").format(*args) | |
| def _finite(fn, args: list[float], description: str) -> float: | |
| """fn(*args), or a student-facing ValueError if it has no finite value.""" | |
| try: | |
| value = float(fn(*args)) | |
| except (ValueError, ZeroDivisionError, OverflowError): | |
| value = math.nan | |
| if not math.isfinite(value): | |
| raise ValueError(f"{description} has no finite value. Try other values.") | |
| return value | |
| def _value(op: str, args: list[float]) -> float: | |
| return _finite(VALUES[op], args, _formula(op, args)) | |
| _UNSUPPORTED = f"Use only numbers, variables, + - * / ^ and the functions {', '.join(FUNCTIONS)}." | |
| class _GraphBuilder: | |
| def __init__(self): | |
| self.inputs = set() | |
| self.consts = {} | |
| self.ops = [] | |
| # (op, args) -> a temporary id, so repeated subexpressions share a node | |
| self.op_ids = {} | |
| def number(self, value: object) -> str: | |
| if type(value) not in (int, float): | |
| raise ValueError(_UNSUPPORTED) | |
| try: | |
| number = float(value) | |
| except OverflowError: | |
| number = math.inf | |
| if not math.isfinite(number): | |
| raise ValueError(f"The number {value} is too large.") | |
| return self.const(repr(number).removesuffix(".0"), number) | |
| def const(self, node_id: str, value: float) -> str: | |
| self.consts[node_id] = value | |
| return node_id | |
| def op(self, node: ast.expr, op: str, *args: str) -> str: | |
| if all(arg in self.consts for arg in args): | |
| label = ast.unparse(node).replace("**", "^") | |
| return self.const(label, _value(op, [self.consts[arg] for arg in args])) | |
| key = (op, args) | |
| if key not in self.op_ids: | |
| self.op_ids[key] = f"#{len(self.ops)}" | |
| self.ops.append({"id": self.op_ids[key], "kind": "op", "op": op, "args": list(args)}) | |
| return self.op_ids[key] | |
| def name(self, name: str) -> str: | |
| if name in CONSTANTS: | |
| return self.const(name, CONSTANTS[name]) | |
| if name in FUNCTIONS: | |
| raise ValueError(f"{name} is a function: write it as {name}(x).") | |
| if re.fullmatch(r"w_?\d+", name): | |
| raise ValueError(f"{name} is the graph's name for an intermediate value: choose another variable name.") | |
| self.inputs.add(name) | |
| return name | |
| def visit(self, node: ast.expr) -> str: | |
| match node: | |
| case ast.Constant(value=value): | |
| return self.number(value) | |
| case ast.Name(id=name): | |
| return self.name(name) | |
| case ast.UnaryOp(op=ast.USub(), operand=ast.Constant(value=value)): | |
| return self.number(-value if type(value) in (int, float) else value) | |
| case ast.UnaryOp(op=ast.USub(), operand=operand): | |
| return self.op(node, "neg", self.visit(operand)) | |
| case ast.UnaryOp(op=ast.UAdd(), operand=operand): | |
| return self.visit(operand) | |
| case ast.BinOp(left=left, op=op, right=right) if type(op) in BINARY_OPS: | |
| left_id = self.visit(left) | |
| return self.op(node, BINARY_OPS[type(op)], left_id, self.visit(right)) | |
| case ast.Call(func=ast.Name(id=name), args=args, keywords=[]): | |
| if name not in FUNCTIONS: | |
| raise ValueError(f"{name} isn't a function here. Use one of: {', '.join(FUNCTIONS)}.") | |
| if len(args) != 1 or isinstance(args[0], ast.Starred): | |
| raise ValueError(f"{name} takes one argument.") | |
| return self.op(node, name, self.visit(args[0])) | |
| raise ValueError(_UNSUPPORTED) | |
| def parse_function(text: str) -> dict: | |
| """The graph of a typed function. Raises ValueError with a student-facing message.""" | |
| text = text.strip() | |
| if not text: | |
| raise ValueError("Enter a function, like x*y + sin(x).") | |
| builder = _GraphBuilder() | |
| try: | |
| # Python's own ^ binds looser than +, so it would read x^2 + 1 as x^(2 + 1) | |
| output = builder.visit(ast.parse(text.replace("^", "**"), mode="eval").body) | |
| except SyntaxError: | |
| raise ValueError("This function can't be read. Check the brackets, and write multiplication with *, like 2*x.") | |
| except (RecursionError, MemoryError): | |
| raise ValueError("This function is too long.") | |
| if not builder.inputs: | |
| raise ValueError("The function must use at least one variable, like x.") | |
| inputs = sorted(builder.inputs) | |
| ids = {name: f"w{i + 1}" for i, name in enumerate(inputs)} | |
| ids.update({op["id"]: f"w{len(inputs) + i + 1}" for i, op in enumerate(builder.ops)}) | |
| ops = [{**op, "id": ids[op["id"]], "args": [ids.get(arg, arg) for arg in op["args"]]} for op in builder.ops] | |
| used = {arg for op in ops for arg in op["args"]} | |
| return { | |
| "inputs": inputs, | |
| "constants": {node_id: value for node_id, value in builder.consts.items() if node_id in used}, | |
| "nodes": [*({"id": ids[name], "kind": "input", "name": name} for name in inputs), *ops], | |
| "output": ids[output], | |
| } | |
| def _partials(node: dict, values: dict, constants: dict) -> list[tuple[str, float]]: | |
| """(arg, local derivative) for each argument that is a node; constants have no derivative.""" | |
| args = [values[arg] for arg in node["args"]] | |
| description = f"The derivative of {_formula(node['op'], args)}" | |
| return [ | |
| (arg, _finite(PARTIALS[node["op"]][i], args, description)) | |
| for i, arg in enumerate(node["args"]) | |
| if arg not in constants | |
| ] | |
| def _step(phase: str, node_id: str, value: float, derivative: float | None, partials: list) -> dict: | |
| step = {"phase": phase, "node": node_id, "value": value, "partials": partials} | |
| if derivative is not None: | |
| step["derivative"] = derivative | |
| return step | |
| def _amount(arg: str, partial: float, amount: float) -> dict: | |
| return {"arg": arg, "partial": partial, "amount": _finite(lambda: amount, [], "The derivative")} | |
| def forward_mode(graph: dict, inputs: dict, seed: str) -> dict: | |
| """A run computing each operation's value and tangent, in one pass, with respect to seed.""" | |
| values = dict(graph["constants"]) | |
| tangents = {} | |
| start = {"values": {}, "derivatives": {}} | |
| steps = [] | |
| for node in graph["nodes"]: | |
| node_id = node["id"] | |
| if node["kind"] == "input": | |
| values[node_id] = start["values"][node_id] = inputs[node["name"]] | |
| tangents[node_id] = start["derivatives"][node_id] = 1.0 if node["name"] == seed else 0.0 | |
| else: | |
| values[node_id] = _value(node["op"], [values[arg] for arg in node["args"]]) | |
| partials = [_amount(arg, p, p * tangents[arg]) for arg, p in _partials(node, values, graph["constants"])] | |
| tangents[node_id] = _finite(sum, [[p["amount"] for p in partials]], "The derivative") | |
| steps.append(_step("forward", node_id, values[node_id], tangents[node_id], partials)) | |
| return {"start": start, "steps": steps} | |
| def reverse_mode(graph: dict, inputs: dict) -> dict: | |
| """A run computing each operation's value going forward, then each adjoint going backward.""" | |
| values = dict(graph["constants"]) | |
| start = {"values": {}, "derivatives": {}} | |
| steps = [] | |
| for node in graph["nodes"]: | |
| if node["kind"] == "input": | |
| values[node["id"]] = start["values"][node["id"]] = inputs[node["name"]] | |
| else: | |
| values[node["id"]] = _value(node["op"], [values[arg] for arg in node["args"]]) | |
| steps.append(_step("forward", node["id"], values[node["id"]], None, [])) | |
| output = graph["output"] | |
| adjoints = {output: 1.0} | |
| steps.append(_step("backward", output, values[output], 1.0, [])) | |
| for node in reversed(graph["nodes"]): | |
| if node["kind"] == "op": | |
| adjoint = adjoints[node["id"]] | |
| partials = [] | |
| for arg, p in _partials(node, values, graph["constants"]): | |
| partials.append(_amount(arg, p, adjoint * p)) | |
| adjoints[arg] = _finite(VALUES["add"], [adjoints.get(arg, 0.0), partials[-1]["amount"]], "The derivative") | |
| steps.append(_step("backward", node["id"], values[node["id"]], adjoint, partials)) | |
| return {"start": start, "steps": steps} | |
| def state_after(run: dict, k: int) -> dict: | |
| """Every value and derivative known after k steps, and the latest step, if any.""" | |
| values = dict(run["start"]["values"]) | |
| derivatives = dict(run["start"]["derivatives"]) | |
| for step in run["steps"][:k]: | |
| values[step["node"]] = step["value"] | |
| if "derivative" in step: | |
| derivatives[step["node"]] = step["derivative"] | |
| if step["phase"] == "backward": | |
| for p in step["partials"]: | |
| derivatives[p["arg"]] = derivatives.get(p["arg"], 0.0) + p["amount"] | |
| state = {"values": values, "derivatives": derivatives} | |
| if k > 0: | |
| state["step"] = run["steps"][k - 1] | |
| return state | |