@tinyvolt/pie 0.0.0-stage → 0.1.0

package/src/lib/nodes/components/NodeCanvas.tsx

@@ -0,0 +1,58 @@
+import { For, createMemo, JSX } from 'solid-js';
+import { NodeStore, type NodeSignal } from '../state/NodeStore';
+import { getNodeDefinition } from '../state/Node';
+import { CenteredGrid } from './CenteredGrid';
+import { useCanvasResize } from '../hooks/useCanvasResize';
+
+interface NodeCanvasProps {
+  store: NodeStore;
+}
+
+export function NodeCanvas(props: NodeCanvasProps): JSX.Element {
+  let svgContainerRef: HTMLDivElement | undefined;
+  useCanvasResize(() => svgContainerRef, props.store);
+
+  // Canvas size in grid units: the viewBox maps one unit to `unit` pixels.
+  const viewWidth = () => props.store.canvasBounds().width / props.store.unit();
+  const viewHeight = () => props.store.canvasBounds().height / props.store.unit();
+
+  // One pass over the store, partitioned into the two draw layers.
+  const layers = createMemo(() => {
+    const geometry: NodeSignal[] = [];
+    const points: NodeSignal[] = [];
+    for (const nodeSignal of Object.values(props.store.nodes)) {
+      if (!nodeSignal) continue;
+      const node = nodeSignal.signal();
+      if (!node || node.type === 'Dummy' || node.hidden) continue;
+      if (node.type === 'Point') points.push(nodeSignal);
+      else geometry.push(nodeSignal);
+    }
+    return { geometry, points };
+  });
+
+  const nodeGroup = (nodeSignal: NodeSignal) => {
+    const node = createMemo(() => nodeSignal?.signal());
+    // The element is rebuilt only when the type changes.
+    const type = createMemo(() => node()?.type);
+    return (
+      <g data-node-id={node()?.id}>
+        {/* The browser shows the `<title>` as a tooltip when hovering any shape in the group. */}
+        <title>{node()?.id}</title>
+        {type() && (getNodeDefinition(type()!)?.render?.(node, props.store.unit) ?? null)}
+      </g>
+    );
+  };
+
+  return (
+    <div class="node-canvas" ref={svgContainerRef}>
+      <svg viewBox={`0 0 ${viewWidth()} ${viewHeight()}`}>
+        <CenteredGrid width={viewWidth()} height={viewHeight()} unit={props.store.unit()}>
+          {/* Step 1: Draw non-`Point` nodes first */}
+          <For each={layers().geometry}>{(nodeSignal) => nodeGroup(nodeSignal)}</For>
+          {/* Step 2: Draw `Point` nodes above geometry */}
+          <For each={layers().points}>{(nodeSignal) => nodeGroup(nodeSignal)}</For>
+        </CenteredGrid>
+      </svg>
+    </div>
+  );
+}

package/src/lib/nodes/execute.ts

@@ -0,0 +1,166 @@
+import { batch, createMemo, createRoot } from 'solid-js';
+import { getFunction, getSideEffect } from './functions/registry';
+import { NodeReconciler, ParameterDefinition } from './functions/Node2NodeFunc';
+import { NodeStore } from './state/NodeStore';
+import { AnyNode, ArrayNode, NodeType, ScalarNode, TextNode, checkColor } from './state/Node';
+import { scalar } from './jax';
+import { leaves, listItemType, listShape, sha256 } from './tensors/tensor';
+import { broadcastItems, broadcastShapes, isBroadcast } from './tensors/broadcast';
+
+/** One argument of a Python call: a number, a string, a handle to an existing node, or a nested list of these. */
+export type Arg = number | string | { id: string } | Arg[];
+
+export interface ExecuteResult {
+  id: string;
+  type: string;
+  /** Only for an Array. */
+  shape?: number[];
+  itemType?: string;
+}
+
+function saveLiteral(store: NodeStore, node: AnyNode): string {
+  createRoot((dispose) => {
+    store.saveNode(node.id, createMemo(() => node), dispose);
+  });
+  return node.id;
+}
+
+/** What an argument is, before any node is created for it. */
+function argType(arg: Arg, store: NodeStore): string {
+  if (Array.isArray(arg)) {
+    listItemType(arg, item => argType(item, store));
+    return 'Array';
+  }
+  if (typeof arg === 'number') return 'Scalar';
+  if (typeof arg === 'string') return 'Text';
+  const node = store.getNodeValue(arg.id);
+  if (!node) throw new Error(`Unknown node '${arg.id}'`);
+  return node.type;
+}
+
+/** The item type and shape of an argument that is an array, before any node is created for it. */
+function arrayOf(arg: Arg, store: NodeStore): { itemType: NodeType; shape: number[] } {
+  if (Array.isArray(arg)) return { itemType: listItemType(arg, item => argType(item, store)), shape: listShape(arg) };
+  return store.getNodeValue((arg as { id: string }).id) as ArrayNode;
+}
+
+/**
+ * A list gets one hidden Array node, reused by every call that passes the same list.
+ * Its items follow the nodes they name.
+ */
+async function resolveList(list: Arg[], store: NodeStore): Promise<string> {
+  const id = `_literal_arr_${await sha256(JSON.stringify(list))}`;
+  if (store.getNodeValue(id)) return id;
+  const shape = listShape(list);
+  const itemType = listItemType(list, item => argType(item, store));
+  const itemIds = await Promise.all(leaves(list).map(item => resolveArg(item, store)));
+  createRoot((dispose) => {
+    const node = createMemo(() => ({ id, type: 'Array', itemType, shape, items: itemIds.map(itemId => store.getNodeValue(itemId)) } as ArrayNode));
+    store.saveNode(id, node, dispose);
+  });
+  return id;
+}
+
+/** A number or string gets one hidden node, reused by every call that passes the same value. */
+async function resolveArg(arg: Arg, store: NodeStore): Promise<string> {
+  if (Array.isArray(arg)) return resolveList(arg, store);
+  if (typeof arg === 'object') return arg.id;
+  const id = typeof arg === 'number' ? `_literal_num_${arg}` : `_literal_str_${arg}`;
+  if (store.getNodeValue(id)) return id;
+  if (typeof arg === 'number') {
+    return saveLiteral(store, { id, type: 'Scalar', value: scalar(arg) } as ScalarNode);
+  }
+  return saveLiteral(store, { id, type: 'Text', value: arg } as TextNode);
+}
+
+/** Checks `out`, the id asked for the output node of a call with arguments `args`. */
+export function checkOut(keyword: string, out: unknown, args: Arg[]) {
+  if (out === undefined) return;
+  if (typeof out !== 'string' || out === '') {
+    throw new Error(`${keyword}: out must be a non-empty string`);
+  }
+  // Ids starting with '_' are kept for the hidden number and text nodes.
+  if (out.startsWith('_')) {
+    throw new Error(`${keyword}: out must not start with '_'`);
+  }
+  // A node that is its own input would update forever.
+  if (leaves(args).some(arg => typeof arg === 'object' && arg.id === out)) {
+    throw new Error(`${keyword}: out='${out}' is also an input`);
+  }
+}
+
+/** The keyword arguments of a Python call. */
+export interface CallOptions {
+  /** The id for the output node. Without it the id is generated. */
+  out?: string;
+  hidden?: boolean;
+  fill?: string;
+  stroke?: string;
+  /** Only for `animatable` side effects: numbers move to their new values instead of jumping. */
+  animate?: boolean;
+}
+
+export async function execute(store: NodeStore, keyword: string, args: Arg[], { out, hidden, fill, stroke, animate }: CallOptions = {}): Promise<ExecuteResult> {
+  const func = getFunction(keyword) ?? getSideEffect(keyword);
+  if (!func) {
+    throw new Error(`Unknown function: ${keyword}`);
+  }
+
+  const total = func.parameters.length;
+  const required = func.parameters.filter(p => p.defaultValue === undefined).length;
+  if (args.length < required || args.length > total) {
+    const names = func.parameters.map((p: ParameterDefinition) => p.name).join(', ');
+    const count = required === total ? `exactly ${total}` : `${required} to ${total}`;
+    throw new Error(`${keyword} needs ${count} argument${total === 1 ? '' : 's'} (${names}), got ${args.length}`);
+  }
+  // Left-out arguments take their defaults. A null default makes no node: `compute` gets undefined.
+  const allArgs = func.parameters.map((p, i) => (i < args.length ? args[i] : p.defaultValue!));
+
+  // Check every type and shape first, so a bad call leaves no literal nodes behind.
+  const broadcastInputShapes: number[][] = [];
+  allArgs.forEach((arg, i) => {
+    if (arg === null) return;
+    const param = func.parameters[i];
+    const types = Array.isArray(param.type) ? param.type : [param.type];
+    let got = argType(arg, store);
+    // TODO: maybe add `broadcastable` flag to each param instead of the entire func?
+    if (isBroadcast(got, param)) {
+      if (func.broadcastable === false) throw new Error(`${keyword} does not take arrays`);
+      const array = arrayOf(arg, store);
+      broadcastInputShapes.push(array.shape);
+      got = `Array of ${array.itemType}`;
+      if (types.includes(array.itemType)) return;
+    }
+    if (!types.includes(got as NodeType)) {
+      throw new Error(`${keyword}: expected ${types.join(' or ')} but got ${got} for argument '${param.name}'`);
+    }
+  });
+  if (broadcastInputShapes.length > 0) broadcastShapes(broadcastInputShapes);
+
+  checkOut(keyword, out, args);
+  if ('run' in func && hidden !== undefined) throw new Error(`${keyword} does not take hidden`);
+  if ('run' in func && fill !== undefined) throw new Error(`${keyword} does not take fill`);
+  if ('run' in func && stroke !== undefined) throw new Error(`${keyword} does not take stroke`);
+  if (animate !== undefined && !('run' in func && func.animatable)) throw new Error(`${keyword} does not take animate`);
+  checkColor(keyword, 'fill', fill);
+  checkColor(keyword, 'stroke', stroke);
+
+  const inputNodeIds = await Promise.all(allArgs.map(arg => (arg === null ? undefined : resolveArg(arg, store))));
+  const inputs = () => inputNodeIds.map(id => (id === undefined ? undefined : store.getNodeValue(id)!) as AnyNode);
+
+  if ('run' in func) {
+    const runs = batch(() => broadcastItems(func.parameters, inputs()).map(itemInputs => func.run(itemInputs, store, animate ?? false)));
+    await Promise.all(runs);
+    return { id: '', type: 'Dummy' };
+  }
+
+  const outputType = broadcastInputShapes.length > 0
+    ? 'Array'
+    : func.resolveOutputType?.(inputs()) ?? func.outputType;
+  const id = out ?? store.nameGenerator.generateAutoName(store.nameGenerator.getPrefixForType(outputType));
+  new NodeReconciler().evaluate(func, inputNodeIds, store, id, hidden, fill, stroke);
+
+  const node = store.getNodeValue(id);
+  if (node?.type === 'Array') return { id, type: node.type, shape: node.shape, itemType: node.itemType };
+  return { id, type: node?.type ?? func.outputType };
+}

package/src/lib/nodes/functions/Node2NodeFunc.ts

@@ -0,0 +1,307 @@
+import { NodeType } from '../state/Node';
+import { batch, createComputed, createMemo, createRoot, onCleanup } from 'solid-js';
+import { AnyNode } from '../state/Node';
+import { NodeStore } from '../state/NodeStore';
+import { computeBlueprint } from '../tensors/broadcast';
+
+export interface ParameterDefinition {
+  name: string;
+  /** A list means any one of these types is accepted. */
+  type: NodeType | NodeType[];
+  /**
+   * Makes the parameter optional: a left-out argument takes this value. `null`: a left-out argument
+   * makes no node, and `compute` gets `undefined`. Optional parameters come after the required ones.
+   */
+  defaultValue?: number | string | null;
+}
+
+export type BlueprintNode = any;
+
+/** Keyword prefix of nodes made by `pie.lift`ed Python functions, followed by the function's name. Not in the registry. */
+export const PYTHON_KEYWORD_PREFIX = '_python:';
+
+/**
+ * How long a reactive recompute error may persist before it is reported.
+ */
+const TRANSIENT_ERROR_GRACE_MS = 500;
+
+export interface Node2NodeFunc {
+  keyword: string;
+  outputType: NodeType;
+  /** For a function whose output type depends on its inputs. Used to name the output. */
+  resolveOutputType?: (inputs: AnyNode[]) => NodeType;
+  parameters: ParameterDefinition[];
+  /** False: passing an array is an error instead of a loop over its items. Default true. */
+  broadcastable?: boolean;
+  compute: (inputs: AnyNode[]) => Record<string, BlueprintNode>;
+}
+
+/** A function that changes the store instead of making a node. It runs once per call, not reactively. */
+export interface SideEffectFunc {
+  keyword: string;
+  parameters: ParameterDefinition[];
+  /** False: passing an array is an error instead of a loop over its items. Default true. */
+  broadcastable?: boolean;
+  /** True: takes `animate`. Default false. */
+  animatable?: boolean;
+  /** The call waits for the returned promise, e.g. until an animation ends. */
+  run: (inputs: AnyNode[], store: NodeStore, animate: boolean) => void | Promise<void>;
+}
+
+type LocalMeta = {
+  globalId: string;
+  disposeNode: () => void;
+  initialized: boolean;
+}
+
+export class NodeReconciler {
+
+  /**
+   * `resolveReference`
+   * When the dep has an allocated id but hasn't been saved yet, we return the original
+   * static value. `getNodeValue` is reactive, so the enclosing memo re-runs automatically
+   * once the dep is saved, upgrading the static reference to the live node.
+   *
+   * `blueprintIndex` maps each blueprint object to its local key, so finding the
+   * key that owns `item` is a hash lookup.
+   */
+  private resolveReference(
+    item: any,
+    localKey: string,
+    localToGlobal: Map<string, LocalMeta>,
+    blueprintIndex: Map<object, string>,
+    store: NodeStore
+  ): { node: any; missing: boolean } | null {
+    if (typeof item !== 'object' || item === null) return null;
+
+    const otherLocalKey = blueprintIndex.get(item);
+    if (otherLocalKey === undefined || otherLocalKey === localKey) return null;
+    const foundMeta = localToGlobal.get(otherLocalKey);
+    if (!foundMeta) return null;
+
+    const globalNode = store.getNodeValue(foundMeta.globalId);
+    if (globalNode) return { node: globalNode, missing: false };
+
+    // Not in store: either not-yet-saved (keep static) or deleted (cascade)
+    return { node: item, missing: foundMeta.initialized };
+  }
+
+  /**
+   * Index every top-level blueprint object by identity, so `resolveReference` is
+   * O(1). First-wins on key order.
+   */
+  private buildBlueprintIndex(blueprint: Record<string, BlueprintNode> | null): Map<object, string> {
+    const index = new Map<object, string>();
+    if (!blueprint) return index;
+    for (const key in blueprint) {
+      const value = blueprint[key];
+      if (value && typeof value === 'object' && !index.has(value)) index.set(value, key);
+    }
+    return index;
+  }
+
+  /**
+   * Walks the node's properties and rewrites any static blueprint references to
+   * reactive lookups via `getNodeValue`. Mutates `transformedData` in place.
+   *
+   * Example: `ellipse_from_foci` yields an ellipse whose `center` is a static
+   * point object that is also a blueprint entry of its own. This rewrites it to
+   * the global reactive node for that point.
+   */
+  private linkBlueprintReferences(
+    localKey: string,
+    transformedData: AnyNode,
+    localToGlobal: Map<string, LocalMeta>,
+    blueprintIndex: Map<object, string>,
+    store: NodeStore
+  ): { hasMissingDependency: boolean } {
+    let hasMissingDependency = false;
+
+    for (const key in transformedData) {
+      const value = transformedData[key as keyof AnyNode];
+      if (typeof value !== 'object' || value === null) continue;
+
+      if (Array.isArray(value)) {
+        (transformedData as any)[key] = (value as any[]).map((item: any) => {
+          const resolved = this.resolveReference(item, localKey, localToGlobal, blueprintIndex, store);
+          if (!resolved) return item;
+          if (resolved.missing) hasMissingDependency = true;
+          return resolved.node;
+        });
+      } else {
+        const resolved = this.resolveReference(value, localKey, localToGlobal, blueprintIndex, store);
+        if (!resolved) continue;
+        if (resolved.missing) hasMissingDependency = true;
+        (transformedData as any)[key] = resolved.node;
+      }
+    }
+
+    return { hasMissingDependency };
+  }
+
+  /**
+   * `hidden`, `fill` and `stroke` are set on every node the call makes: the main node and the others in the blueprint.
+   * An `undefined` input id is a left-out optional parameter: `compute` gets `undefined` for it.
+   */
+  evaluate(func: Node2NodeFunc, inputNodeIds: (string | undefined)[], store: NodeStore, assignedId?: string, hidden = false, fill?: string, stroke?: string) {
+    createRoot((disposeRoot) => {
+      const localToGlobal = new Map<string, LocalMeta>();
+
+      // A compute error during recompute keeps the previous blueprint instead of
+      // escaping mid-propagation. It is reported only if it persists past the
+      // grace window.
+      let firstRunDone = false;
+      let graceTimer: ReturnType<typeof setTimeout> | undefined;
+      let lastErrorMessage = '';
+      onCleanup(() => { if (graceTimer !== undefined) clearTimeout(graceTimer); });
+
+      // The inputs the current blueprint was computed from. Set together with the
+      // blueprint, so a kept previous blueprint keeps its own inputs.
+      let blueprintInputs: AnyNode[] = [];
+
+      const blueprintMemo = createMemo<Record<string, BlueprintNode> | null>((prev) => {
+        const inputs = inputNodeIds.map(id => id === undefined ? undefined : store.getNodeValue(id));
+        if (inputs.some((node, i) => inputNodeIds[i] !== undefined && (!node || node.type === 'Dummy'))) return null;
+        try {
+          const blueprint = computeBlueprint(func, inputs as AnyNode[]);
+          blueprintInputs = inputs as AnyNode[];
+          if (graceTimer !== undefined) {
+            clearTimeout(graceTimer);
+            graceTimer = undefined;
+          }
+          return blueprint;
+        } catch (err) {
+          // First run happens at call time — let the caller report it.
+          if (!firstRunDone) throw err;
+          lastErrorMessage = err instanceof Error ? err.message : String(err);
+          if (graceTimer === undefined) {
+            graceTimer = setTimeout(() => {
+              graceTimer = undefined;
+              console.error(`${func.keyword}: ${lastErrorMessage}`);
+            }, TRANSIENT_ERROR_GRACE_MS);
+          }
+          return prev;
+        }
+      }, null);
+      firstRunDone = true;
+
+      this.reconcile(blueprintMemo, () => blueprintInputs, localToGlobal, func, assignedId, hidden, fill, stroke, store, disposeRoot);
+    }) // outer `createRoot` ends here
+  }
+
+  private reconcile(
+    blueprintAccessor: () => Record<string, BlueprintNode> | null,
+    inputsAccessor: () => AnyNode[],
+    localToGlobal: Map<string, LocalMeta>,
+    func: Node2NodeFunc,
+    assignedId: string | undefined,
+    hidden: boolean,
+    fill: string | undefined,
+    stroke: string | undefined,
+    store: NodeStore,
+    disposeRoot: () => void
+  ) {
+    // One identity index per blueprint value, shared by every node memo below.
+    const blueprintIndex = createMemo(() => this.buildBlueprintIndex(blueprintAccessor()));
+
+    createComputed(() => {
+      const blueprint = blueprintAccessor();
+      const currentKeys = blueprint ? new Set<string>(Object.keys(blueprint)) : new Set<string>();
+
+      // Phase A: create missing nodes
+      if (blueprint) {
+        const newlyAddedKeys: string[] = [];
+        for (const localKey of currentKeys) {
+          if (!localToGlobal.has(localKey)) {
+            const isMain = localKey === 'main';
+            const nodeData = blueprint[localKey];
+            const autoNamePrefix = store.nameGenerator.getPrefixForType(nodeData.type);
+
+            const globalId = isMain && assignedId ? assignedId : store.nameGenerator.generateAutoName(autoNamePrefix);
+
+            localToGlobal.set(localKey, { globalId, disposeNode: () => { }, initialized: false });
+            newlyAddedKeys.push(localKey);
+          }
+        }
+
+        // Phase A2: build reactive roots and save each node to the global store.
+        batch(() => {
+          for (const localKey of newlyAddedKeys) {
+            const meta = localToGlobal.get(localKey)!;
+            const { globalId } = meta;
+            const isMain = localKey === 'main';
+
+            createRoot((disposeNode) => {
+              meta.disposeNode = disposeNode;
+
+              const nodeSignal = createMemo(() => {
+                const latestBlueprint = blueprintAccessor();
+                if (!latestBlueprint) return { id: globalId, type: 'Dummy' } as unknown as AnyNode;
+                const data = latestBlueprint[localKey];
+                // Guard against the reconciler not yet having disposed this node
+                // after the blueprint shrank.
+                if (!data) return { id: globalId, type: 'Dummy' } as unknown as AnyNode;
+
+                const transformedData = {
+                  ...data,
+                  id: globalId,
+                  hidden,
+                  fill,
+                  stroke,
+                  __graph_deps: {
+                    funcKeyword: func.keyword,
+                    inputNodes: inputsAccessor(),
+                    blueprintKey: localKey,
+                    ...(localKey !== 'main' ? { parentId: localToGlobal.get('main')?.globalId } : {})
+                  }
+                } as AnyNode;
+
+                const { hasMissingDependency } = this.linkBlueprintReferences(
+                  localKey, transformedData, localToGlobal, blueprintIndex(), store
+                );
+
+                if (hasMissingDependency) {
+                  queueMicrotask(() => {
+                    if (store.getNodeValue(globalId)) store.removeNode(globalId);
+                  });
+                  return { id: globalId, type: 'Dummy' } as unknown as AnyNode;
+                }
+
+                return transformedData;
+              });
+
+              const customDispose = () => {
+                disposeNode();
+                if (isMain) {
+                  for (const [otherLocalKey, otherMeta] of localToGlobal.entries()) {
+                    if (otherLocalKey !== localKey && store.getNodeValue(otherMeta.globalId)) {
+                      store.removeNode(otherMeta.globalId);
+                    }
+                  }
+                  disposeRoot();
+                }
+              };
+
+              store.saveNode(globalId, nodeSignal, customDispose);
+              meta.initialized = true;
+            });
+          }
+        }); // `batch` ends here
+      } // `if (blueprint)` ends here
+
+      // Phase B: clean up removed nodes
+      for (const [localKey, meta] of localToGlobal.entries()) {
+        if (!currentKeys.has(localKey)) {
+          if (store.getNodeValue(meta.globalId)) {
+            store.removeNode(meta.globalId);
+          }
+          localToGlobal.delete(localKey);
+        }
+      }
+
+      if (!blueprint) {
+        disposeRoot();
+      }
+    }); // `createComputed` ends here
+  }
+}

package/src/lib/nodes/functions/array.ts

@@ -0,0 +1,36 @@
+import { Node2NodeFunc } from './Node2NodeFunc';
+import { ArrayNode, TextNode } from '../state/Node';
+
+/**
+ * The items of an array that a slice picks, for Python's `ps[0:2]` or `ps[:, 1]`.
+ * Hidden: `pie._index` is blocked, so it is only reached through `Array.__getitem__`.
+ * @param array - Array to slice
+ * @param key - JSON list with one entry per leading dimension: an index (the dimension is dropped) or a list of indices (the dimension is kept). Dimensions past the key are kept whole.
+ * @returns Array node holding the existing item nodes
+ * @auxiliaryNodes None
+ */
+export const Index: Node2NodeFunc = {
+  keyword: '_index', outputType: 'Array',
+  parameters: [{ name: 'array', type: 'Array' }, { name: 'key', type: 'Text' }],
+  compute(inputs) {
+    const array = inputs[0] as ArrayNode;
+    const key: (number | number[])[] = JSON.parse((inputs[1] as TextNode).value);
+    const picks = array.shape.map((size, d) => {
+      const entry = key[d];
+      if (entry === undefined) return [...Array(size).keys()];
+      return typeof entry === 'number' ? [entry] : entry;
+    });
+    const shape = picks.filter((_, d) => typeof key[d] !== 'number').map(pick => pick.length);
+
+    // Row-major positions in `array.items` of every picked item.
+    let positions = [0];
+    picks.forEach((pick, d) => {
+      positions = positions.flatMap(position => pick.map(i => position * array.shape[d] + i));
+    });
+
+    return { main: { type: 'Array', itemType: array.itemType, shape, items: positions.map(p => array.items[p]) } };
+  },
+};
+
+/** Docs category for the functions defined in this file. */
+export const category = 'Arrays';

package/src/lib/nodes/functions/autodiff.ts

@@ -0,0 +1,64 @@
+import { Node2NodeFunc, PYTHON_KEYWORD_PREFIX } from './Node2NodeFunc';
+import { getFunction } from './registry';
+import { AnyNode, ScalarNode } from '../state/Node';
+import { ref, Dtype } from '../jax';
+import { Graph, replay, gradient } from '~/lib/autodiff';
+import { computeBlueprint } from '../tensors/broadcast';
+
+/** The number fields of each node type `partial` can differentiate with respect to. */
+const NUMBER_FIELDS: Record<string, string[]> = { Scalar: ['value'], Point: ['x', 'y'] };
+
+/** Nodes as an autodiff graph: each node carries its input nodes in `__graph_deps`. */
+const nodeGraph: Graph<AnyNode> = {
+  id: (node) => node.id,
+  inputs: (node) => node.__graph_deps?.inputNodes,
+  recompute(node, inputs) {
+    const deps = node.__graph_deps!;
+    if (deps.funcKeyword.startsWith(PYTHON_KEYWORD_PREFIX)) {
+      throw new Error(`partial can't differentiate through Python function '${deps.funcKeyword.slice(PYTHON_KEYWORD_PREFIX.length)}'`);
+    }
+    const func = getFunction(deps.funcKeyword);
+    if (!func) throw new Error(`Unknown function: ${deps.funcKeyword}`);
+    const key = deps.blueprintKey ?? 'main';
+    const data = computeBlueprint(func, inputs)[key];
+    if (!data) throw new Error(`${deps.funcKeyword} gave no '${key}' output`);
+    return { ...data, id: node.id, __graph_deps: { ...deps, inputNodes: inputs } };
+  },
+};
+
+/**
+ * Partial derivative of a scalar with respect to a scalar or a point.
+ * Composes: the target may itself be a `partial`.
+ * @param target - Scalar to differentiate
+ * @param var - Scalar or Point to differentiate with respect to. Its own inputs are ignored.
+ * @returns Scalar (var is a Scalar) or PointGrad (var is a Point)
+ * @auxiliaryNodes None
+ */
+export const Partial: Node2NodeFunc = {
+  keyword: 'partial', 
+  outputType: 'Scalar',
+  resolveOutputType: (inputs) => (inputs[1].type === 'Point' ? 'PointGrad' : 'Scalar'),
+  broadcastable: false,
+  parameters: [{ name: 'target', type: 'Scalar' }, { name: 'var', type: ['Scalar', 'Point'] }],
+  compute(inputs) {
+    const target = inputs[0] as ScalarNode;
+    const v = inputs[1];
+    if (v.id.startsWith('_literal_')) {
+      throw new Error('partial: var must be a node, not a number');
+    }
+
+    const fields = NUMBER_FIELDS[v.type];
+    // The variable with new values and no inputs, so the replay stops at it.
+    const varAt = (values: Dtype[]) => ({
+      ...v, ...Object.fromEntries(fields.map((field, i) => [field, values[i]])), __graph_deps: undefined,
+    }) as AnyNode;
+    const f = (...values: Dtype[]) => ref((replay(nodeGraph, target, v.id, varAt(values)) as ScalarNode).value);
+    const grads = gradient(f, fields.map(field => ref((v as any)[field])));
+
+    if (v.type === 'Scalar') return { main: { type: 'Scalar', value: grads[0] } };
+    return { main: { type: 'PointGrad', point: v, target, dx: grads[0], dy: grads[1] } };
+  },
+};
+
+/** Docs category for the functions defined in this file. */
+export const category = 'Autodiff';