@jax-js/jax 0.0.1

package/dist/index.js

@@ -0,0 +1,3708 @@
+import {
+  AluExp,
+  AluVar,
+  DEBUG,
+  DType,
+  FpHash,
+  Kernel,
+  Reduction,
+  ShapeTracker,
+  __callDispose,
+  __export,
+  __using,
+  accessorAluExp,
+  accessorGlobal,
+  deepEqual,
+  devices,
+  getBackend,
+  init,
+  invertPermutation,
+  isFloatDtype,
+  isPermutation,
+  partitionList,
+  prod,
+  range,
+  recursiveFlatten,
+  rep,
+  runWithCache,
+  setDevice,
+  toposort,
+  unravelAlu,
+  unzip2,
+  zip
+} from "./chunk-B2GFURUN.js";
+
+// src/pprint.ts
+var PPrint = class _PPrint {
+  constructor(indents, lines) {
+    this.indents = indents;
+    this.lines = lines;
+  }
+  /** Add a fixed amount of indentation to each line. */
+  indent(spaces) {
+    return new _PPrint(
+      this.indents.map((i) => i + spaces),
+      this.lines
+    );
+  }
+  /** Concatenate two or more pretty-printed expressions. */
+  concat(...items) {
+    return new _PPrint(
+      (this.indents ?? []).concat(...items.map((i) => i.indents)),
+      (this.lines ?? []).concat(...items.map((i) => i.lines))
+    );
+  }
+  /** Stack one block to the right of another one, sharing 1 common line. */
+  stack(other) {
+    if (!other.lines.length) return this;
+    if (!this.lines.length) return other;
+    const indent = this.indents[this.indents.length - 1];
+    const s = this.lines[this.lines.length - 1];
+    const indentedBlock = other.indent(indent + s.length);
+    return new _PPrint(
+      this.indents.concat(indentedBlock.indents.slice(1)),
+      this.lines.slice(0, -1).concat(
+        s + " ".repeat(other.indents[0]) + other.lines[0],
+        ...indentedBlock.lines.slice(1)
+      )
+    );
+  }
+  /** Combine this block of lines into a formatted string. */
+  toString() {
+    return this.lines.map((line, i) => " ".repeat(this.indents[i]) + line).join("\n");
+  }
+  static pp(s) {
+    const lines = s.toString().split("\n");
+    return new _PPrint(Array(lines.length).fill(0), lines);
+  }
+};
+
+// src/tree.ts
+var tree_exports = {};
+__export(tree_exports, {
+  JsTreeDef: () => JsTreeDef,
+  NodeType: () => NodeType,
+  flatten: () => flatten,
+  leaves: () => leaves,
+  structure: () => structure,
+  unflatten: () => unflatten
+});
+var NodeType = /* @__PURE__ */ ((NodeType2) => {
+  NodeType2["Array"] = "Array";
+  NodeType2["Object"] = "Object";
+  NodeType2["Leaf"] = "Leaf";
+  return NodeType2;
+})(NodeType || {});
+var JsTreeDef = class _JsTreeDef {
+  constructor(nodeType, nodeMetadata, childTreedefs) {
+    this.nodeType = nodeType;
+    this.nodeMetadata = nodeMetadata;
+    this.childTreedefs = childTreedefs;
+  }
+  static leaf = new _JsTreeDef("Leaf" /* Leaf */, null, []);
+  /** Returns a string representation of this tree definition. */
+  toString(root = true) {
+    if (root) {
+      return "JsTreeDef(" + this.toString(false) + ")";
+    }
+    switch (this.nodeType) {
+      case "Leaf" /* Leaf */:
+        return "*";
+      case "Array" /* Array */:
+        return `[${this.childTreedefs.map((x) => x.toString(false)).join(", ")}]`;
+      case "Object" /* Object */: {
+        const parts = [];
+        for (let i = 0; i < this.childTreedefs.length; i++) {
+          parts.push(
+            `${quoteObjectKey(this.nodeMetadata[i])}: ${this.childTreedefs[i].toString(false)}`
+          );
+        }
+        return `{${parts.join(", ")}}`;
+      }
+    }
+  }
+  /** Compare this tree definition with another. */
+  equals(other) {
+    return this.nodeType === other.nodeType && deepEqual(this.nodeMetadata, other.nodeMetadata) && this.childTreedefs.length === other.childTreedefs.length && this.childTreedefs.every((x, i) => x.equals(other.childTreedefs[i]));
+  }
+};
+function quoteObjectKey(key) {
+  if (/^[a-zA-Z_$][a-zA-Z0-9_$]*$/.test(key)) {
+    return key;
+  }
+  return JSON.stringify(key);
+}
+function flatten(tree) {
+  const leaves2 = [];
+  const treedef = _flatten(tree, leaves2);
+  return [leaves2, treedef];
+}
+function _flatten(tree, leaves2) {
+  if (Array.isArray(tree)) {
+    const childTrees = tree.map((c) => _flatten(c, leaves2));
+    return new JsTreeDef("Array" /* Array */, null, childTrees);
+  } else if (typeof tree === "object" && tree !== null && tree.constructor === Object) {
+    const [keys, values] = unzip2(Object.entries(tree));
+    const childTrees = values.map((c) => _flatten(c, leaves2));
+    return new JsTreeDef("Object" /* Object */, keys, childTrees);
+  } else {
+    leaves2.push(tree);
+    return JsTreeDef.leaf;
+  }
+}
+function leaves(tree) {
+  return flatten(tree)[0];
+}
+function structure(tree) {
+  return flatten(tree)[1];
+}
+function unflatten(treedef, leaves2) {
+  return _unflatten(treedef, leaves2[Symbol.iterator]());
+}
+function _unflatten(treedef, leaves2) {
+  switch (treedef.nodeType) {
+    case "Leaf" /* Leaf */: {
+      const { value, done } = leaves2.next();
+      if (done) {
+        throw new TypeError("Ran out of leaves while unflattening JsTree");
+      }
+      return value;
+    }
+    case "Array" /* Array */:
+      return treedef.childTreedefs.map((c) => _unflatten(c, leaves2));
+    case "Object" /* Object */: {
+      const obj = {};
+      for (let i = 0; i < treedef.childTreedefs.length; i++) {
+        obj[treedef.nodeMetadata[i]] = _unflatten(
+          treedef.childTreedefs[i],
+          leaves2
+        );
+      }
+      return obj;
+    }
+  }
+}
+
+// src/frontend/core.ts
+function add(x, y) {
+  return bind1("add" /* Add */, [x, y]);
+}
+function mul(x, y) {
+  return bind1("mul" /* Mul */, [x, y]);
+}
+function idiv(x, y) {
+  return bind1("idiv" /* Idiv */, [x, y]);
+}
+function neg(x) {
+  return bind1("neg" /* Neg */, [x]);
+}
+function reciprocal(x) {
+  return bind1("reciprocal" /* Reciprocal */, [x]);
+}
+function sin(x) {
+  return bind1("sin" /* Sin */, [x]);
+}
+function cos(x) {
+  return bind1("cos" /* Cos */, [x]);
+}
+function exp(x) {
+  return bind1("exp" /* Exp */, [x]);
+}
+function log(x) {
+  return bind1("log" /* Log */, [x]);
+}
+function min(x, y) {
+  return bind1("min" /* Min */, [x, y]);
+}
+function max(x, y) {
+  return bind1("max" /* Max */, [x, y]);
+}
+function compare(x, y, op) {
+  return bind1("compare" /* Compare */, [x, y], { op });
+}
+function greater(x, y) {
+  return compare(x, y, "greater" /* Greater */);
+}
+function less(x, y) {
+  return compare(x, y, "less" /* Less */);
+}
+function equal(x, y) {
+  return compare(x, y, "equal" /* Equal */);
+}
+function notEqual(x, y) {
+  return compare(x, y, "not_equal" /* NotEqual */);
+}
+function greaterEqual(x, y) {
+  return compare(x, y, "greater_equal" /* GreaterEqual */);
+}
+function lessEqual(x, y) {
+  return compare(x, y, "less_equal" /* LessEqual */);
+}
+function where(cond, x, y) {
+  return bind1("where" /* Where */, [cond, x, y]);
+}
+function transpose(x, perm) {
+  perm = perm ?? range(ndim(x)).reverse();
+  return bind1("transpose" /* Transpose */, [x], { perm });
+}
+function broadcast(x, shape2, axis) {
+  return bind1("broadcast" /* Broadcast */, [x], { shape: shape2, axis });
+}
+function reshape(x, shape2) {
+  if (typeof shape2 === "number") shape2 = [shape2];
+  const originalShape = getShape(x);
+  const autoIdx = shape2.indexOf(-1);
+  if (autoIdx !== -1) {
+    const remaining = prod(originalShape) / -prod(shape2);
+    if (!Number.isInteger(remaining) || remaining < 0) {
+      throw new TypeError(
+        `Invalid reshape: ${JSON.stringify(originalShape)} -> ${JSON.stringify(shape2)}`
+      );
+    }
+    shape2 = shape2.toSpliced(autoIdx, 1, remaining);
+  }
+  if (prod(originalShape) !== prod(shape2)) {
+    throw new TypeError(
+      `Invalid reshape: ${JSON.stringify(originalShape)} -> ${JSON.stringify(shape2)}`
+    );
+  }
+  return bind1("reshape" /* Reshape */, [x], { shape: shape2 });
+}
+function flip(x, axis) {
+  return bind1("flip" /* Flip */, [x], { axis });
+}
+function reduceSum(x, axis) {
+  if (axis === void 0) {
+    if (x instanceof Tracer) {
+      axis = range(x.shape.length);
+    } else {
+      axis = [];
+    }
+  }
+  if (typeof axis === "number") {
+    axis = [axis];
+  }
+  return bind1("reduce_sum" /* ReduceSum */, [x], { axis });
+}
+function bind1(prim, args, params = {}) {
+  const [results] = bind(prim, args, params);
+  return results;
+}
+var traceStack = [];
+var dynamicTrace = null;
+function newMain(traceType, globalData = null) {
+  const level = traceStack.length;
+  const main = { level, traceType, globalData };
+  traceStack.push(main);
+  return Object.assign(main, {
+    [Symbol.dispose]() {
+      traceStack.pop();
+    }
+  });
+}
+function newDynamic(main) {
+  const prevDynamicTrace = dynamicTrace;
+  dynamicTrace = main;
+  return {
+    [Symbol.dispose]() {
+      dynamicTrace = prevDynamicTrace;
+    }
+  };
+}
+var Trace = class {
+  constructor(main) {
+    this.main = main;
+  }
+};
+var Tracer = class {
+  _trace;
+  constructor(trace) {
+    this._trace = trace;
+  }
+  get shape() {
+    return this.aval.shape;
+  }
+  get dtype() {
+    return this.aval.dtype;
+  }
+  get ndim() {
+    return this.shape.length;
+  }
+  fullLower() {
+    return this;
+  }
+  // These types aren't technically correct since they don't account for the
+  // fact that tracers can be lifted to different levels. But they simplify the
+  // API visible to users.
+  neg() {
+    return neg(this);
+  }
+  add(other) {
+    return add(this, other);
+  }
+  mul(other) {
+    return mul(this, other);
+  }
+  greater(other) {
+    return greater(this, other);
+  }
+  less(other) {
+    return less(this, other);
+  }
+  equal(other) {
+    return equal(this, other);
+  }
+  notEqual(other) {
+    return notEqual(this, other);
+  }
+  greaterEqual(other) {
+    return greaterEqual(this, other);
+  }
+  lessEqual(other) {
+    return lessEqual(this, other);
+  }
+  sum(axis) {
+    return reduceSum(this, axis);
+  }
+  transpose(perm) {
+    return transpose(this, perm);
+  }
+  reshape(shape2) {
+    return reshape(this, shape2);
+  }
+  // Below this line are composite operations built from primitives.
+  /** Subtract an array from this one. */
+  sub(other) {
+    return this.add(neg(other));
+  }
+  /** Divide an array by this one. */
+  div(other) {
+    if (this.dtype === "int32" /* Int32 */) {
+      return idiv(this, other);
+    }
+    return this.mul(reciprocal(other));
+  }
+  /** Return specified diagonals. See `numpy.diagonal` for full docs. */
+  diagonal(offset = 0, axis1 = 0, axis2 = 1) {
+    if (!Number.isInteger(offset))
+      throw new TypeError(`offset must be an integer, got ${offset}`);
+    if (axis1 === axis2)
+      throw new TypeError("axis1 and axis2 must not be equal");
+    throw new Error("diagonal not implemented");
+  }
+  /** Flatten the array without changing its data. */
+  flatten() {
+    return this.reshape(-1);
+  }
+  /** Flatten the array without changing its data. */
+  ravel() {
+    return this.reshape(-1);
+  }
+};
+function ndim(x) {
+  if (x instanceof Tracer) {
+    return x.shape.length;
+  } else {
+    return 0;
+  }
+}
+function getShape(x) {
+  return x instanceof Tracer ? x.shape : [];
+}
+var ShapedArray = class _ShapedArray {
+  constructor(shape2, dtype) {
+    this.shape = shape2;
+    this.dtype = dtype;
+  }
+  static fromAval(aval) {
+    return new _ShapedArray(aval.shape, aval.dtype);
+  }
+  get ndim() {
+    return this.shape.length;
+  }
+  strShort() {
+    return `${this.dtype}[${this.shape.join(",")}]`;
+  }
+  equals(other) {
+    return this === other || this.constructor === other.constructor && this.ndim === other.ndim && this.shape.every((d, i) => d === other.shape[i]);
+  }
+};
+function getAval(x) {
+  if (x instanceof Tracer) {
+    return x.aval;
+  } else if (typeof x === "boolean" || typeof x === "number") {
+    return new ShapedArray(
+      [],
+      typeof x === "boolean" ? "bool" /* Bool */ : "float32" /* Float32 */
+    );
+  } else {
+    throw new TypeError(`Unknown value: ${x}`);
+  }
+}
+function bind(prim, args, params = {}) {
+  const topTrace = findTopTrace(args);
+  const tracers = args.map((arg) => fullRaise(topTrace, arg));
+  const outs = topTrace.processPrimitive(prim, tracers, params);
+  if (DEBUG >= 5) {
+    console.info(
+      `processing rule for ${prim} on ${tracers.map((x) => x.toString())} and got ${outs.map((x) => x.toString())}`
+    );
+  }
+  return outs.map((out) => out.fullLower());
+}
+function findTopTrace(xs) {
+  let topMain = traceStack[0];
+  for (const x of xs) {
+    if (x instanceof Tracer && x._trace.main.level > topMain.level) {
+      topMain = x._trace.main;
+    }
+  }
+  if (dynamicTrace && dynamicTrace.level > topMain.level) {
+    topMain = dynamicTrace;
+  }
+  return new topMain.traceType(topMain);
+}
+function fullRaise(trace, val) {
+  if (!(val instanceof Tracer)) {
+    return trace.pure(val);
+  }
+  const level = trace.main.level;
+  if (Object.is(val._trace.main, trace.main)) {
+    return val;
+  } else if (val._trace.main.level < level) {
+    return trace.lift(val);
+  } else if (val._trace.main.level > level) {
+    throw new Error(
+      `Can't lift Tracer level ${val._trace.main.level} to level ${level}`
+    );
+  } else {
+    throw new Error(`Different traces at same level: ${val._trace}, ${trace}.`);
+  }
+}
+var TreeMismatchError = class extends TypeError {
+  constructor(where3, left, right) {
+    super(`Mismatched tree structures in ${where3}: ${left} != ${right}`);
+  }
+};
+function flattenFun(f, inTree) {
+  const store = { value: void 0 };
+  const flatFun = (...argsFlat) => {
+    const pytreeArgs = unflatten(inTree, argsFlat);
+    const out = f(...pytreeArgs);
+    const [outFlat, outTree] = flatten(out);
+    store.value = outTree;
+    return outFlat;
+  };
+  return [flatFun, store];
+}
+var UseAfterFreeError = class extends ReferenceError {
+  constructor(tracer) {
+    super(
+      `Referenced tracer ${tracer.toString()} freed, please use .ref move semantics`
+    );
+  }
+};
+
+// src/frontend/jit.ts
+var JitProgram = class {
+  constructor(backend, steps, inputs, outputs) {
+    this.backend = backend;
+    this.steps = steps;
+    this.inputs = inputs;
+    this.outputs = outputs;
+  }
+  /** Execute the JitProgram with the given inputs. */
+  execute(inputs) {
+    const scope = /* @__PURE__ */ new Map();
+    if (inputs.length !== this.inputs.length) {
+      throw new TypeError(
+        `Expected ${this.inputs.length} inputs, got ${inputs.length}`
+      );
+    }
+    for (const [i, id] of this.inputs.entries()) {
+      scope.set(id, inputs[i]);
+    }
+    const pending = [];
+    for (const step of this.steps) {
+      switch (step.type) {
+        case "execute": {
+          const inputs2 = step.inputs.map((id) => scope.get(id));
+          const outputs = step.outputs.map((id) => scope.get(id));
+          if (inputs2.some((s) => s === void 0) || outputs.some((s) => s === void 0)) {
+            throw new Error(`internal: JitProgram scope undefined`);
+          }
+          pending.push(
+            new PendingExecute(this.backend, step.kernel, inputs2, outputs)
+          );
+          break;
+        }
+        case "const":
+          scope.set(step.output, step.slot);
+          break;
+        case "malloc": {
+          const slot = this.backend.malloc(4 * step.size);
+          scope.set(step.output, slot);
+          break;
+        }
+        case "free": {
+          const slot = scope.get(step.input);
+          this.backend.decRef(slot);
+          scope.delete(step.input);
+          break;
+        }
+      }
+    }
+    return {
+      outputs: this.outputs.map((id) => scope.get(id)),
+      pending
+    };
+  }
+};
+var JitProgramBuilder = class {
+  backend;
+  #nextId;
+  steps;
+  constructor(backend, nargs) {
+    this.backend = backend;
+    this.#nextId = nargs;
+    this.steps = [];
+  }
+  pushConst(slot) {
+    const id = this.#nextId++;
+    this.steps.push({
+      type: "const",
+      slot,
+      output: id
+    });
+    return id;
+  }
+  pushLit(lit) {
+    const kernel = new Kernel(
+      0,
+      prod(lit.aval.shape),
+      AluExp.const(lit.dtype, lit.value)
+    );
+    return this.pushKernel(kernel, []);
+  }
+  pushKernel(kernel, inputs) {
+    const id = this.#nextId++;
+    this.steps.push({
+      type: "malloc",
+      size: kernel.size,
+      output: id
+    });
+    this.steps.push({
+      type: "execute",
+      kernel,
+      inputs,
+      outputs: [id]
+    });
+    return id;
+  }
+  insertFreeSteps(outputIds) {
+    const ids = this.steps.filter((s) => s.type === "malloc").map((s) => s.output);
+    for (const id of ids) {
+      if (outputIds.includes(id)) continue;
+      const lastUsage = this.steps.findLastIndex(
+        (s) => s.type === "execute" && (s.outputs.includes(id) || s.inputs.includes(id)) || s.type === "malloc" && s.output === id
+      );
+      this.steps.splice(lastUsage + 1, 0, {
+        type: "free",
+        input: id
+      });
+    }
+  }
+  pushFree(id) {
+    this.steps.push({
+      type: "free",
+      input: id
+    });
+  }
+};
+var jitCompileCache = /* @__PURE__ */ new Map();
+function jitCompile(backend, jaxpr, consts) {
+  if (jaxpr.inBinders.length < consts.length) {
+    throw new TypeError(
+      `Jaxpr has ${jaxpr.inBinders.length} inputs, but ${consts.length} consts were provided`
+    );
+  }
+  for (let i = 0; i < consts.length; i++) {
+    if (consts[i].device !== backend.type) {
+      throw new TypeError(
+        `Const ${i} has device ${consts[i].device}, but expected ${backend.type}`
+      );
+    }
+  }
+  const cacheKey = backend.type + FpHash.hash(jaxpr, ...consts.map((c) => c.id));
+  const cached = jitCompileCache.get(cacheKey);
+  if (cached) return cached;
+  if (DEBUG >= 1) {
+    console.info("=========== JIT Compile ===========\n" + jaxpr.toString());
+  }
+  jaxpr = jaxpr.flatten().simplify();
+  const nargs = jaxpr.inBinders.length - consts.length;
+  const builder = new JitProgramBuilder(backend, nargs);
+  const blackNodes = splitGraphDataflow(backend, jaxpr);
+  const ctx = /* @__PURE__ */ new Map();
+  for (let i = 0; i < consts.length; i++) {
+    const v = jaxpr.inBinders[i];
+    const slot = consts[i]._realizeSource();
+    ctx.set(v, { type: "imm", arg: builder.pushConst(slot) });
+  }
+  for (let i = 0; i < nargs; i++) {
+    const v = jaxpr.inBinders[consts.length + i];
+    ctx.set(v, { type: "imm", arg: i });
+  }
+  for (let i = 0; i < jaxpr.eqns.length; i++) {
+    const eqn = jaxpr.eqns[i];
+    const inputExps = [];
+    const inputAvals = [];
+    const inputArgs = [];
+    for (const input of eqn.inputs) {
+      if (input instanceof Var) {
+        const jitValue = ctx.get(input);
+        if (jitValue.type === "exp") {
+          const gidMap = /* @__PURE__ */ new Map();
+          for (const [gid, jitId] of jitValue.args.entries()) {
+            let newGid = inputArgs.indexOf(jitId);
+            if (newGid === -1) {
+              newGid = inputArgs.length;
+              inputArgs.push(jitId);
+            }
+            gidMap.set(gid, newGid);
+          }
+          inputExps.push(jitValue.exp.reindexGids(gidMap));
+        } else if (jitValue.type === "imm") {
+          let gid = inputArgs.indexOf(jitValue.arg);
+          if (gid === -1) {
+            gid = inputArgs.length;
+            inputArgs.push(jitValue.arg);
+          }
+          const st = ShapeTracker.fromShape(input.aval.shape);
+          const indices = unravelAlu(st.shape, AluVar.gidx);
+          inputExps.push(AluExp.globalView(input.aval.dtype, gid, st, indices));
+        } else {
+          jitValue;
+        }
+        inputAvals.push(input.aval);
+      } else if (input instanceof Lit) {
+        inputExps.push(AluExp.const(input.dtype, input.value));
+        inputAvals.push(input.aval);
+      } else {
+        throw new TypeError(`Unexpected input in Jaxpr: ${input}`);
+      }
+    }
+    const nargs2 = inputArgs.length;
+    const rule = jitRules[eqn.primitive];
+    if (!rule)
+      throw new TypeError(`JIT not implemented for primitive ${eqn.primitive}`);
+    const kernel = rule(nargs2, inputExps, inputAvals, eqn.params);
+    const outVar = eqn.outBinders[0];
+    if (kernel.reduction || blackNodes.has(outVar)) {
+      const outId = builder.pushKernel(kernel, inputArgs);
+      ctx.set(outVar, { type: "imm", arg: outId });
+    } else {
+      ctx.set(outVar, { type: "exp", exp: kernel.exp, args: inputArgs });
+    }
+  }
+  const outputIds = [];
+  for (const out of jaxpr.outs) {
+    if (out instanceof Var) {
+      const jitValue = ctx.get(out);
+      if (jitValue.type !== "imm")
+        throw new Error("internal: Expected imm, since outs are black nodes");
+      outputIds.push(jitValue.arg);
+    } else if (out instanceof Lit) {
+      outputIds.push(builder.pushLit(out));
+    } else {
+      out;
+    }
+  }
+  builder.insertFreeSteps(outputIds);
+  const jp = new JitProgram(backend, builder.steps, range(0, nargs), outputIds);
+  jitCompileCache.set(cacheKey, jp);
+  return jp;
+}
+function broadcastedJit(fn) {
+  return (nargs, exps, avals, params) => {
+    const newShape = avals.map((aval) => aval.shape).reduce(generalBroadcast);
+    exps = exps.map(
+      (exp4) => exp4.rewrite((exp5) => {
+        if (exp5.op === "GlobalView" /* GlobalView */) {
+          let [gid, st] = exp5.arg;
+          if (!deepEqual(st.shape, newShape)) {
+            st = st.broadcast(
+              newShape,
+              range(newShape.length - st.shape.length)
+            );
+            const indices = unravelAlu(st.shape, AluVar.gidx);
+            return AluExp.globalView(exp5.dtype, gid, st, indices);
+          }
+        }
+      })
+    );
+    const exp3 = fn(exps, params);
+    return new Kernel(nargs, prod(newShape), exp3);
+  };
+}
+function reshapeJit(fn) {
+  return (nargs, [a], [as], params) => {
+    a = a.rewrite((exp3) => {
+      if (exp3.op === "GlobalView" /* GlobalView */) {
+        const [gid, st] = exp3.arg;
+        const newSt = fn(st, params);
+        const indices = unravelAlu(newSt.shape, AluVar.gidx);
+        return AluExp.globalView(exp3.dtype, gid, newSt, indices);
+      }
+    });
+    return new Kernel(nargs, prod(as.shape), a);
+  };
+}
+var jitRules = {
+  ["add" /* Add */]: broadcastedJit(([a, b]) => AluExp.add(a, b)),
+  ["mul" /* Mul */]: broadcastedJit(([a, b]) => AluExp.mul(a, b)),
+  ["idiv" /* Idiv */]: broadcastedJit(([a, b]) => AluExp.idiv(a, b)),
+  ["neg" /* Neg */]: broadcastedJit(
+    ([a]) => AluExp.sub(AluExp.const(a.dtype, 0), a)
+  ),
+  ["reciprocal" /* Reciprocal */]: broadcastedJit(([a]) => AluExp.reciprocal(a)),
+  ["sin" /* Sin */]: broadcastedJit(([a]) => AluExp.sin(a)),
+  ["cos" /* Cos */]: broadcastedJit(([a]) => AluExp.cos(a)),
+  ["exp" /* Exp */]: broadcastedJit(([a]) => AluExp.exp(a)),
+  ["log" /* Log */]: broadcastedJit(([a]) => AluExp.log(a)),
+  ["min" /* Min */]: broadcastedJit(([a, b]) => AluExp.min(a, b)),
+  ["max" /* Max */]: broadcastedJit(([a, b]) => AluExp.max(a, b)),
+  ["reduce_sum" /* ReduceSum */](nargs, [a], [as], { axis }) {
+    const keptAxes = [];
+    const shiftedAxes = [];
+    const newShape = [];
+    for (let i = 0; i < as.shape.length; i++) {
+      if (axis.includes(i)) shiftedAxes.push(i);
+      else {
+        keptAxes.push(i);
+        newShape.push(as.shape[i]);
+      }
+    }
+    const size2 = prod(newShape);
+    const reductionSize = prod(shiftedAxes.map((ax) => as.shape[ax]));
+    newShape.push(reductionSize);
+    a = a.rewrite((exp3) => {
+      if (exp3.op === "GlobalView" /* GlobalView */) {
+        const [gid, st] = exp3.arg;
+        const newSt = st.permute(keptAxes.concat(shiftedAxes)).reshape(newShape);
+        const indices = unravelAlu(newShape.slice(0, -1), AluVar.gidx);
+        indices.push(AluVar.ridx);
+        return AluExp.globalView(exp3.dtype, gid, newSt, indices);
+      }
+    });
+    const reduction = new Reduction(a.dtype, "Add" /* Add */, reductionSize);
+    return new Kernel(nargs, size2, a, reduction);
+  },
+  ["compare" /* Compare */]: broadcastedJit(([a, b], { op }) => {
+    return aluCompare(a, b, op);
+  }),
+  ["where" /* Where */]: broadcastedJit(([cond, a, b]) => AluExp.where(cond, a, b)),
+  ["transpose" /* Transpose */]: reshapeJit(
+    (st, { perm }) => {
+      return st.permute(perm);
+    }
+  ),
+  ["broadcast" /* Broadcast */]: reshapeJit(
+    (st, { shape: shape2, axis }) => {
+      return st.broadcast(shape2, axis);
+    }
+  ),
+  ["reshape" /* Reshape */]: reshapeJit(
+    (st, { shape: shape2 }) => {
+      return st.reshape(shape2);
+    }
+  ),
+  ["flip" /* Flip */]: reshapeJit(
+    (st, { axis }) => {
+      const arg = rep(st.shape.length, false);
+      for (const ax of axis) arg[ax] = true;
+      return st.flip(arg);
+    }
+  )
+};
+function splitGraphDataflow(backend, jaxpr) {
+  const varToEqn = /* @__PURE__ */ new Map();
+  for (let i = 0; i < jaxpr.eqns.length; i++) {
+    const eqn = jaxpr.eqns[i];
+    for (const v of eqn.outBinders) {
+      if (v instanceof Var) varToEqn.set(v, i);
+    }
+  }
+  const blackNodes = /* @__PURE__ */ new Set();
+  const p1NextBlack = /* @__PURE__ */ new Map();
+  for (const v of jaxpr.outs) {
+    if (v instanceof Var) {
+      blackNodes.add(v);
+      p1NextBlack.set(v, v);
+    }
+  }
+  for (let i = jaxpr.eqns.length - 1; i >= 0; i--) {
+    const eqn = jaxpr.eqns[i];
+    if (eqn.primitive === "reduce_sum" /* ReduceSum */ || eqn.outBinders.some((v) => blackNodes.has(v))) {
+      for (const v of eqn.outBinders) {
+        blackNodes.add(v);
+        p1NextBlack.set(v, v);
+      }
+      continue;
+    }
+    const reach = /* @__PURE__ */ new Set();
+    for (let j = i + 1; j < jaxpr.eqns.length; j++) {
+      for (const v of jaxpr.eqns[j].inputs) {
+        if (v instanceof Var && eqn.outBinders.includes(v)) {
+          for (const o of jaxpr.eqns[j].outBinders) {
+            const u = p1NextBlack.get(o);
+            if (u) reach.add(u);
+          }
+        }
+      }
+    }
+    if (reach.size === 1) {
+      const b = reach.values().next().value;
+      for (const v of eqn.outBinders) p1NextBlack.set(v, b);
+    } else if (reach.size > 1) {
+      for (const v of eqn.outBinders) {
+        blackNodes.add(v);
+        p1NextBlack.set(v, v);
+      }
+    }
+  }
+  const p2Deps = /* @__PURE__ */ new Map();
+  for (const v of jaxpr.inBinders) {
+    p2Deps.set(v, /* @__PURE__ */ new Set([v]));
+  }
+  let p2idx = 0;
+  while (p2idx < jaxpr.eqns.length) {
+    const eqn = jaxpr.eqns[p2idx++];
+    const deps = [];
+    if (eqn.outBinders.some((v) => blackNodes.has(v))) {
+      continue;
+    }
+    for (const input of eqn.inputs) {
+      if (input instanceof Var) {
+        if (blackNodes.has(input)) deps.push(/* @__PURE__ */ new Set([input]));
+        else deps.push(p2Deps.get(input));
+      } else {
+        deps.push(/* @__PURE__ */ new Set());
+      }
+    }
+    const depCounter = /* @__PURE__ */ new Map();
+    for (const depSet of deps) {
+      for (const dep of depSet) {
+        depCounter.set(dep, (depCounter.get(dep) ?? 0) + 1);
+      }
+    }
+    if (depCounter.size > backend.maxArgs) {
+      let maxUniqueDeps = 0;
+      let assocInput = -1;
+      for (let i = 0; i < eqn.inputs.length; i++) {
+        const input = eqn.inputs[i];
+        if (input instanceof Var && varToEqn.has(input)) {
+          let uniqueDeps = 0;
+          for (const dep of deps[i]) {
+            if (depCounter.get(dep) === 1) uniqueDeps++;
+          }
+          if (uniqueDeps > maxUniqueDeps) {
+            maxUniqueDeps = uniqueDeps;
+            assocInput = i;
+          }
+        }
+      }
+      if (assocInput === -1) {
+        throw new Error(
+          `internal: maxArgs, no input found to mark as black in Jaxpr equation ${eqn}`
+        );
+      }
+      const assocVar = eqn.inputs[assocInput];
+      p2idx = varToEqn.get(assocVar);
+      for (const out of jaxpr.eqns[p2idx].outBinders) {
+        blackNodes.add(out);
+      }
+    } else {
+      const s = new Set(depCounter.keys());
+      for (const out of eqn.outBinders) p2Deps.set(out, s);
+    }
+  }
+  return blackNodes;
+}
+
+// src/frontend/array.ts
+var JsArray = globalThis.Array;
+var inlineArrayLimit = 1e3;
+var fudgeArray = pureArray;
+var PendingExecute = class {
+  // since this could be held by multiple arrays, cancel when it hits 0
+  constructor(backend, kernel, inputs, outputs) {
+    this.backend = backend;
+    this.kernel = kernel;
+    this.inputs = inputs;
+    this.outputs = outputs;
+    for (const slot of inputs) this.backend.incRef(slot);
+  }
+  prepared = null;
+  submitted = false;
+  #promise = null;
+  // for prepare
+  #rc = 1;
+  // Change the reference count of the PendingExecute object.
+  // Used when copying the object to a new Array, or disposing an array.
+  updateRc(delta) {
+    if (this.#rc <= 0) throw new Error("internal: PendingExecute used rc<=0");
+    this.#rc += delta;
+    if (this.#rc <= 0 && !this.submitted) {
+      for (const slot of this.inputs) this.backend.decRef(slot);
+    }
+  }
+  async prepare() {
+    if (this.prepared) return;
+    if (this.#promise) {
+      await this.#promise;
+      return;
+    }
+    this.#promise = (async () => {
+      this.prepared = await this.backend.prepare(this.kernel);
+    })();
+    await this.#promise;
+  }
+  prepareSync() {
+    if (this.prepared) return;
+    this.prepared = this.backend.prepareSync(this.kernel);
+  }
+  submit() {
+    if (this.submitted) return;
+    if (this.#rc <= 0) throw new Error("internal: PendingExecute used rc<=0");
+    if (!this.prepared) throw new Error("internal: Not prepared yet");
+    this.submitted = true;
+    this.backend.dispatch(this.prepared, this.inputs, this.outputs);
+    for (const slot of this.inputs) this.backend.decRef(slot);
+  }
+};
+var Array3 = class _Array extends Tracer {
+  static #nextId = 1001;
+  // For unique hashing where needed.
+  id;
+  #dtype;
+  #source;
+  #st;
+  #backend;
+  #rc;
+  // reference count for this specific Array object
+  #pendingSet;
+  // only if source is `Slot`
+  constructor(source, st, dtype, backend, pending = null) {
+    super(baseArrayTrace);
+    this.id = _Array.#nextId++;
+    this.#dtype = dtype;
+    this.#source = source;
+    this.#st = st;
+    this.#backend = backend;
+    this.#rc = 1;
+    this.#pendingSet = new Set(pending);
+    if (!(source instanceof AluExp)) {
+      backend.incRef(source);
+    }
+  }
+  get aval() {
+    return new ShapedArray(this.#st.shape, this.#dtype);
+  }
+  /** Return a simple string representation of the array's dimensions. */
+  toString() {
+    return `Array:${this.#dtype}[${this.shape.join(",")}]`;
+  }
+  get device() {
+    return this.#backend.type;
+  }
+  #check() {
+    if (this.#rc <= 0) throw new UseAfterFreeError(this);
+  }
+  get ref() {
+    this.#check();
+    this.#rc++;
+    return this;
+  }
+  dispose() {
+    this.#check();
+    if (this.#rc-- === 0) {
+      for (const exe of this.#pending) exe.updateRc(-1);
+      if (typeof this.#source === "number") {
+        this.#backend.decRef(this.#source);
+      }
+    }
+  }
+  /** Get the pending executes as a list, trimming if already submitted. */
+  get #pending() {
+    if (!this.#pendingSet) return [];
+    for (const p of this.#pendingSet) {
+      if (p.submitted) this.#pendingSet.delete(p);
+    }
+    if (this.#pendingSet.size === 0) {
+      this.#pendingSet = null;
+      return [];
+    } else {
+      return [...this.#pendingSet];
+    }
+  }
+  /**
+   * Convert this array into a primitive value.
+   *
+   * This only works for scalars (0-dimensional arrays). It lets you get values
+   * "out" of the JAX system. For instance, if `x = np.array(5)`, then you can
+   * evaluate `x + 1` and `x ** 2` to get `6` and `25`, respectively.
+   *
+   * This method is also called for `==` equality.
+   */
+  [Symbol.toPrimitive]() {
+    if (this.ndim === 0) {
+      return this.dataSync()[0];
+    } else {
+      throw new Error(
+        `Cannot convert non-scalar array to primitive: ${this.toString()}`
+      );
+    }
+  }
+  #reshape(st) {
+    this.#check();
+    const pending = this.#pending;
+    for (const exe of pending) exe.updateRc(1);
+    const ar = new _Array(this.#source, st, this.#dtype, this.#backend, pending);
+    this.dispose();
+    return ar;
+  }
+  /** Move axes to the rightmost dimension of the shape. */
+  #moveAxesDown(axis) {
+    this.#check();
+    if (axis.length === 0) return this.reshape(this.shape.concat(1));
+    const newShape = [];
+    const keptAxes = [];
+    const shiftedAxes = [];
+    for (let i = 0; i < this.#st.shape.length; i++) {
+      if (axis.includes(i)) {
+        shiftedAxes.push(i);
+      } else {
+        keptAxes.push(i);
+        newShape.push(this.#st.shape[i]);
+      }
+    }
+    newShape.push(-1);
+    return this.#transpose(keptAxes.concat(shiftedAxes)).reshape(newShape);
+  }
+  #transpose(perm) {
+    this.#check();
+    if (!isPermutation(perm, this.ndim))
+      throw new Error(`Invalid perm for transpose: ${JSON.stringify(perm)}`);
+    return this.#reshape(this.#st.permute(perm));
+  }
+  #unary(op) {
+    this.#check();
+    if (this.#source instanceof AluExp) {
+      const exp4 = new AluExp(op, this.#dtype, [this.#source]);
+      return new _Array(exp4, this.#st, this.#dtype, this.#backend);
+    }
+    const indices = unravelAlu(this.#st.shape, AluVar.gidx);
+    const exp3 = new AluExp(op, this.#dtype, [
+      AluExp.globalView(this.#dtype, 0, this.#st, indices)
+    ]);
+    const kernel = new Kernel(1, this.#st.size, exp3);
+    const output = this.#backend.malloc(kernel.size * 4);
+    const pending = [...this.#pending];
+    for (const exe of pending) exe.updateRc(1);
+    pending.push(
+      new PendingExecute(this.#backend, kernel, [this.#source], [output])
+    );
+    this.dispose();
+    return new _Array(
+      output,
+      ShapeTracker.fromShape(this.shape),
+      this.#dtype,
+      this.#backend,
+      pending
+    );
+  }
+  #binary(op, other) {
+    const custom = (src) => new AluExp(op, this.#dtype, src);
+    return _Array.#naryCustom(op, custom, [this, other]);
+  }
+  static #naryCustom(name, custom, arrays, dtypeOverride, dtypeOutput) {
+    const n = arrays.length;
+    const backend = arrays[0].#backend;
+    if (n === 0) throw new TypeError(`No inputs for ${name}`);
+    for (const ar of arrays) ar.#check();
+    let dtype;
+    for (let i = 0; i < n; i++) {
+      if (dtypeOverride?.[i]) {
+        if (arrays[i].#dtype !== dtypeOverride[i]) {
+          throw new TypeError(
+            `Wrong dtype in ${name}: expected ${dtypeOverride[i]}, got ${arrays[i].#dtype}`
+          );
+        }
+      } else {
+        if (!dtype) dtype = arrays[i].#dtype;
+        else if (arrays[i].#dtype !== dtype) {
+          throw new TypeError(
+            `Dtype mismatch in ${name}: ${dtype} vs ${arrays[i].#dtype}`
+          );
+        }
+      }
+      if (arrays[i].#backend !== backend) {
+        throw new TypeError(
+          `Backend mismatch in ${name}: ${backend.type} vs ${arrays[i].#backend.type}`
+        );
+      }
+    }
+    dtypeOutput ??= dtype;
+    if (!dtypeOutput) throw new TypeError("nary operation with no dtype");
+    const newShape = arrays.map((a) => a.shape).reduce(generalBroadcast);
+    arrays = arrays.map((ar) => {
+      if (deepEqual(ar.shape, newShape)) return ar;
+      return ar.#reshape(
+        ar.#st.broadcast(newShape, range(newShape.length - ar.ndim))
+      );
+    });
+    if (arrays.every((ar) => ar.#source instanceof AluExp)) {
+      if (arrays.every((ar) => deepEqual(ar.#st, arrays[0].#st))) {
+        const exp5 = custom(arrays.map((ar) => ar.#source));
+        return new _Array(exp5, arrays[0].#st, exp5.dtype, backend);
+      }
+      const exp4 = custom(
+        arrays.map((ar) => {
+          const src2 = ar.#source;
+          if (ar.#st.contiguous) return src2;
+          return accessorAluExp(src2, ar.#st, unravelAlu(newShape, AluVar.idx));
+        })
+      );
+      const st = ShapeTracker.fromShape(newShape);
+      return new _Array(exp4, st, exp4.dtype, backend);
+    }
+    const inputs = [];
+    const src = [];
+    for (const ar of arrays) {
+      const indices = unravelAlu(newShape, AluVar.gidx);
+      if (ar.#source instanceof AluExp) {
+        src.push(accessorAluExp(ar.#source, ar.#st, indices));
+      } else {
+        let gid = inputs.indexOf(ar.#source);
+        if (gid === -1) {
+          gid = inputs.length;
+          inputs.push(ar.#source);
+        }
+        src.push(AluExp.globalView(ar.#dtype, gid, ar.#st, indices));
+      }
+    }
+    const exp3 = custom(src);
+    const kernel = new Kernel(inputs.length, arrays[0].#st.size, exp3);
+    const output = backend.malloc(kernel.size * 4);
+    const pending = [...arrays.flatMap((ar) => ar.#pending)];
+    for (const exe of pending) exe.updateRc(1);
+    pending.push(new PendingExecute(backend, kernel, inputs, [output]));
+    for (const ar of arrays) ar.dispose();
+    return new _Array(
+      output,
+      ShapeTracker.fromShape(newShape),
+      dtypeOutput,
+      backend,
+      pending
+    );
+  }
+  /** Reduce the last dimension of the array by an operation. */
+  #reduce(op) {
+    this.#check();
+    if (this.ndim === 0) throw new Error("Cannot reduce a scalar");
+    const shape2 = this.shape;
+    const reduction = new Reduction(this.#dtype, op, shape2[shape2.length - 1]);
+    const newShape = shape2.slice(0, -1);
+    const newSize = prod(newShape);
+    const indices = [...unravelAlu(newShape, AluVar.gidx), AluVar.ridx];
+    const [index, valid] = this.#st.toAluExp(indices);
+    let exp3;
+    const inputs = [];
+    if (this.#source instanceof AluExp) {
+      exp3 = AluExp.where(
+        valid,
+        this.#source.substitute({ idx: index }),
+        AluExp.f32(0)
+      );
+    } else {
+      inputs.push(this.#source);
+      exp3 = AluExp.where(
+        valid,
+        AluExp.globalIndex("float32" /* Float32 */, 0, index),
+        AluExp.f32(0)
+      );
+    }
+    const kernel = new Kernel(inputs.length, newSize, exp3, reduction);
+    const output = this.#backend.malloc(kernel.size * 4);
+    const pending = [...this.#pending];
+    for (const exe of pending) exe.updateRc(1);
+    pending.push(new PendingExecute(this.#backend, kernel, inputs, [output]));
+    this.dispose();
+    return new _Array(
+      output,
+      ShapeTracker.fromShape(newShape),
+      this.#dtype,
+      this.#backend,
+      pending
+    );
+  }
+  /**
+   * Normalizes this array into one backed by a `Slot`.
+   *
+   * This mutates the array in-place, turning it into an equivalent array whose
+   * source is actual, contiguous data on device.
+   *
+   * Calling this twice is a no-op.
+   */
+  #realize() {
+    this.#check();
+    const indices = unravelAlu(this.#st.shape, AluVar.gidx);
+    if (this.#source instanceof AluExp) {
+      const exp3 = accessorAluExp(this.#source, this.#st, indices);
+      const kernel = new Kernel(0, this.#st.size, exp3);
+      const output = this.#backend.malloc(kernel.size * 4);
+      const pendingItem = new PendingExecute(
+        this.#backend,
+        kernel,
+        [],
+        [output]
+      );
+      this.#source = output;
+      this.#st = ShapeTracker.fromShape(this.shape);
+      this.#pendingSet = /* @__PURE__ */ new Set([pendingItem]);
+    } else {
+      if (this.#st.contiguous) return;
+      const exp3 = accessorGlobal(this.#dtype, 0, this.#st, indices);
+      const kernel = new Kernel(1, this.#st.size, exp3);
+      const output = this.#backend.malloc(kernel.size * 4);
+      const pendingItem = new PendingExecute(
+        this.#backend,
+        kernel,
+        [this.#source],
+        [output]
+      );
+      this.#source = output;
+      this.#st = ShapeTracker.fromShape(this.shape);
+      this.#pendingSet ??= /* @__PURE__ */ new Set();
+      this.#pendingSet.add(pendingItem);
+    }
+  }
+  #dataInline() {
+    this.#check();
+    const exp3 = this.#source;
+    const ar = new _Array(exp3, this.#st, this.dtype, getBackend("cpu"));
+    this.dispose();
+    return ar.dataSync();
+  }
+  /** Realize the array and return it as data. */
+  async data() {
+    if (this.#source instanceof AluExp && prod(this.shape) < inlineArrayLimit && this.device !== "cpu") {
+      return this.#dataInline();
+    }
+    this.#realize();
+    const pending = this.#pending;
+    if (pending) {
+      await Promise.all(pending.map((p) => p.prepare()));
+      for (const p of pending) p.submit();
+    }
+    const buf = await this.#backend.read(this.#source);
+    this.dispose();
+    return this.dtype === "float32" /* Float32 */ ? new Float32Array(buf) : new Int32Array(buf);
+  }
+  /** Wait for this array to be placed on the backend, if needed. */
+  async wait() {
+    this.#check();
+    if (this.#source instanceof AluExp) return;
+    const pending = this.#pending;
+    if (pending) {
+      await Promise.all(pending.map((p) => p.prepare()));
+      for (const p of pending) p.submit();
+    }
+    await this.#backend.read(this.#source, 0, 0);
+    this.dispose();
+  }
+  /**
+   * Realize the array and return it as data. This is a sync variant and not
+   * recommended for performance reasons, as it will block rendering.
+   */
+  dataSync() {
+    if (this.#source instanceof AluExp && prod(this.shape) < inlineArrayLimit && this.device !== "cpu") {
+      return this.#dataInline();
+    }
+    this.#realize();
+    for (const p of this.#pending) {
+      p.prepareSync();
+      p.submit();
+    }
+    const buf = this.#backend.readSync(this.#source);
+    this.dispose();
+    return this.dtype === "float32" /* Float32 */ ? new Float32Array(buf) : new Int32Array(buf);
+  }
+  /** Convert this array into a JavaScript object (blocking). */
+  js() {
+    return dataToJs(this.dtype, this.dataSync(), this.shape);
+  }
+  /** Convert this array into a JavaScript object, asynchronously. */
+  async jsAsync() {
+    return dataToJs(this.dtype, await this.data(), this.shape);
+  }
+  /** @private Internal plumbing method for Array / Tracer ops. */
+  static _implRules() {
+    return {
+      ["add" /* Add */]([x, y]) {
+        return [x.#binary("Add" /* Add */, y)];
+      },
+      ["mul" /* Mul */]([x, y]) {
+        return [x.#binary("Mul" /* Mul */, y)];
+      },
+      ["idiv" /* Idiv */]([x, y]) {
+        return [x.#binary("Idiv" /* Idiv */, y)];
+      },
+      ["neg" /* Neg */]([x]) {
+        return [zerosLike(x).#binary("Sub" /* Sub */, x)];
+      },
+      ["reciprocal" /* Reciprocal */]([x]) {
+        return [x.#unary("Reciprocal" /* Reciprocal */)];
+      },
+      ["sin" /* Sin */]([x]) {
+        return [x.#unary("Sin" /* Sin */)];
+      },
+      ["cos" /* Cos */]([x]) {
+        return [x.#unary("Cos" /* Cos */)];
+      },
+      ["exp" /* Exp */]([x]) {
+        return [x.#unary("Exp" /* Exp */)];
+      },
+      ["log" /* Log */]([x]) {
+        return [x.#unary("Log" /* Log */)];
+      },
+      ["min" /* Min */]([x, y]) {
+        return [x.#binary("Min" /* Min */, y)];
+      },
+      ["max" /* Max */]([x, y]) {
+        return [x.#binary("Max" /* Max */, y)];
+      },
+      ["reduce_sum" /* ReduceSum */]([x], { axis }) {
+        if (axis.length === 0) return [x];
+        return [x.#moveAxesDown(axis).#reduce("Add" /* Add */)];
+      },
+      ["compare" /* Compare */]([x, y], { op }) {
+        const custom = ([x2, y2]) => aluCompare(x2, y2, op);
+        return [_Array.#naryCustom("compare", custom, [x, y], [], "bool" /* Bool */)];
+      },
+      ["where" /* Where */]([cond, x, y]) {
+        const custom = ([cond2, x2, y2]) => AluExp.where(cond2, x2, y2);
+        return [_Array.#naryCustom("where", custom, [cond, x, y], ["bool" /* Bool */])];
+      },
+      ["transpose" /* Transpose */]([x], { perm }) {
+        return [x.#transpose(perm)];
+      },
+      ["broadcast" /* Broadcast */]([x], { shape: shape2, axis }) {
+        return [x.#reshape(x.#st.broadcast(shape2, axis))];
+      },
+      ["reshape" /* Reshape */]([x], { shape: shape2 }) {
+        return [x.#reshape(x.#st.reshape(shape2))];
+      },
+      ["flip" /* Flip */]([x], { axis }) {
+        const arg = rep(x.ndim, false);
+        for (const ax of axis) arg[ax] = true;
+        return [x.#reshape(x.#st.flip(arg))];
+      },
+      ["jit_call" /* JitCall */](args, { jaxpr, numConsts }) {
+        if (jaxpr.inBinders.length !== args.length) {
+          throw new Error(
+            `jit_call expects ${jaxpr.inBinders.length} args, got ${args.length}`
+          );
+        }
+        const backend = getBackend();
+        const consts = args.slice(0, numConsts);
+        const tracers = args.slice(numConsts);
+        const jp = jitCompile(backend, jaxpr, consts);
+        const { outputs, pending } = jp.execute(
+          tracers.map((x) => x._realizeSource())
+        );
+        const prevPending = args.flatMap((x) => x.#pending);
+        for (const exe of prevPending) exe.updateRc(1);
+        pending.splice(0, 0, ...prevPending);
+        args.forEach((x) => x.dispose());
+        return outputs.map((source, i) => {
+          return new _Array(
+            source,
+            ShapeTracker.fromShape(jaxpr.outs[i].aval.shape),
+            jaxpr.outs[i].aval.dtype,
+            backend,
+            pending
+          );
+        });
+      }
+    };
+  }
+  // Internal methods, not public API. Do not use.
+  _realizeSource() {
+    this.#realize();
+    return this.#source;
+  }
+};
+function scalar(value, { dtype, device } = {}) {
+  if (typeof value === "number") {
+    dtype ??= "float32" /* Float32 */;
+    if (!["float32" /* Float32 */, "int32" /* Int32 */].includes(dtype))
+      throw new TypeError(`Mismatched dtype for scalar ${value}`);
+  } else if (typeof value === "boolean") {
+    dtype ??= "bool" /* Bool */;
+    if (!["float32" /* Float32 */, "int32" /* Int32 */, "bool" /* Bool */].includes(dtype))
+      throw new TypeError(`Mismatched dtype for scalar ${value}`);
+  } else {
+    throw new TypeError(`Invalid type for scalar ${value}`);
+  }
+  return new Array3(
+    AluExp.const(dtype, value),
+    ShapeTracker.fromShape([]),
+    dtype,
+    getBackend(device)
+  );
+}
+function array(values, { shape: shape2, dtype, device } = {}) {
+  if (values instanceof Array3) {
+    if (shape2 && !deepEqual(values.shape, shape2)) {
+      values = values.reshape(shape2);
+    }
+    if (dtype && values.dtype !== dtype) {
+      throw new Error("array astype not implemented yet");
+    }
+    return values;
+  } else if (values instanceof Float32Array || values instanceof Int32Array) {
+    return arrayFromData(values, shape2 ?? [values.length], {
+      dtype,
+      device
+    });
+  } else {
+    if (!shape2) {
+      shape2 = [];
+      let cur = values;
+      while (JsArray.isArray(cur)) {
+        shape2.push(cur.length);
+        cur = cur[0];
+      }
+    }
+    const size2 = prod(shape2);
+    const flat = recursiveFlatten(values);
+    if (flat.length !== size2) {
+      throw new Error(
+        `Jagged shape: ${JSON.stringify(shape2)} vs ${flat.length}`
+      );
+    }
+    if (size2 === 0) return zeros(shape2, { dtype, device });
+    if (typeof flat[0] === "boolean") {
+      dtype = dtype ?? "bool" /* Bool */;
+      const data = new Int32Array(flat.map((x) => x ? 1 : 0));
+      return arrayFromData(data, shape2, { dtype, device });
+    } else {
+      dtype = dtype ?? "float32" /* Float32 */;
+      const data = new Float32Array(flat);
+      return arrayFromData(data, shape2, { dtype, device });
+    }
+  }
+}
+function arrayFromData(data, shape2, { dtype, device } = {}) {
+  if (data.length < inlineArrayLimit) {
+    let allEqual = true;
+    for (let i = 1; i < data.length; i++) {
+      if (data[i] !== data[0]) {
+        allEqual = false;
+        break;
+      }
+    }
+    if (allEqual) {
+      return full(shape2, data[0], { dtype, device });
+    }
+  }
+  const backend = getBackend(device);
+  if (data instanceof Float32Array) {
+    if (dtype && dtype !== "float32" /* Float32 */) {
+      throw new Error("Float32Array must have float32 type");
+    }
+    const slot = backend.malloc(data.byteLength, data.buffer);
+    return new Array3(
+      slot,
+      ShapeTracker.fromShape(shape2),
+      "float32" /* Float32 */,
+      backend
+    );
+  } else if (data instanceof Int32Array) {
+    if (dtype && dtype !== "int32" /* Int32 */ && dtype !== "bool" /* Bool */) {
+      throw new Error("Int32Array must have int32 or bool type");
+    }
+    const slot = backend.malloc(data.byteLength, data.buffer);
+    return new Array3(
+      slot,
+      ShapeTracker.fromShape(shape2),
+      dtype ?? "int32" /* Int32 */,
+      backend
+    );
+  } else {
+    throw new Error("Unsupported data type");
+  }
+}
+function dataToJs(dtype, data, shape2) {
+  if (shape2.length === 0) {
+    return dtype === "bool" /* Bool */ ? Boolean(data[0]) : data[0];
+  }
+  const [first, ...rest] = shape2;
+  const restSize = prod(rest);
+  const ret = [];
+  for (let i = 0; i < first; i++) {
+    const subarray = data.slice(i * restSize, (i + 1) * restSize);
+    ret.push(dataToJs(dtype, subarray, rest));
+  }
+  return ret;
+}
+function pureArray(x) {
+  if (x instanceof Tracer) {
+    return x;
+  } else {
+    return scalar(x);
+  }
+}
+var EvalTrace = class extends Trace {
+  // No boxing in Tracers needed.
+  pure = (x) => pureArray(x);
+  lift = (x) => x;
+  processPrimitive(primitive, tracers, params) {
+    return implRules[primitive](tracers, params);
+  }
+};
+var baseArrayTrace = new EvalTrace(newMain(EvalTrace, null));
+var implRules = Array3._implRules();
+function zerosLike(val) {
+  const aval = getAval(val);
+  return zeros(aval.shape, { dtype: aval.dtype });
+}
+function zeros(shape2, { dtype, device } = {}) {
+  return full(shape2, 0, { dtype, device });
+}
+function ones(shape2, { dtype, device } = {}) {
+  return full(shape2, 1, { dtype, device });
+}
+function full(shape2, fillValue, { dtype, device } = {}) {
+  let source;
+  if (typeof fillValue === "number") {
+    dtype = dtype ?? "float32" /* Float32 */;
+    source = AluExp.const(dtype, fillValue);
+  } else if (typeof fillValue === "boolean") {
+    dtype = dtype ?? "bool" /* Bool */;
+    source = AluExp.const(dtype, fillValue ? 1 : 0);
+  } else if (fillValue instanceof Array3) {
+    throw new Error("numpy.full() with array argument not implemented yet");
+  } else {
+    throw new TypeError(`Invalid type for full: ${fillValue}`);
+  }
+  return new Array3(
+    source,
+    ShapeTracker.fromShape(shape2),
+    dtype ?? "float32" /* Float32 */,
+    getBackend(device)
+  );
+}
+function eye(numRows, numCols, { dtype, device } = {}) {
+  numCols = numCols ?? numRows;
+  dtype = dtype ?? "float32" /* Float32 */;
+  if (numCols < numRows) {
+    const arr = eye(numCols, numRows, { dtype, device });
+    return arr.transpose();
+  }
+  if (numRows === 0) {
+    return zeros([0, numCols], { dtype, device });
+  }
+  const exp3 = AluExp.cmplt(
+    AluExp.mod(AluVar.idx, AluExp.i32(numCols + 1)),
+    AluExp.i32(1)
+  );
+  return new Array3(
+    AluExp.cast(dtype, exp3),
+    ShapeTracker.fromShape([numRows, numCols]),
+    dtype,
+    getBackend(device)
+  );
+}
+function identity(n, { dtype, device } = {}) {
+  return eye(n, n, { dtype, device });
+}
+function arange(start, stop, step = 1, { dtype, device } = {}) {
+  dtype = dtype ?? "int32" /* Int32 */;
+  if (stop === void 0) {
+    stop = start;
+    start = 0;
+  }
+  if (step === 0) {
+    throw new RangeError(
+      `Invalid step for arange: ${step}. Step must be non-zero.`
+    );
+  }
+  const size2 = Math.max(0, Math.ceil((stop - start) / step));
+  if (size2 === 0) {
+    return zeros([0], { dtype, device });
+  }
+  const exp3 = AluExp.add(
+    AluExp.const(dtype, start),
+    AluExp.mul(AluExp.cast(dtype, AluVar.idx), AluExp.const(dtype, step))
+  );
+  const st = ShapeTracker.fromShape([size2]);
+  return new Array3(exp3, st, dtype, getBackend(device));
+}
+function linspace(start, stop, num = 50, endpoint = true, { dtype, device } = {}) {
+  dtype = dtype ?? "float32" /* Float32 */;
+  if (num < 0 || !Number.isInteger(num)) {
+    throw new RangeError(
+      `Invalid num for linspace: ${num}. Must be non-negative integer.`
+    );
+  } else if (num === 0) {
+    return zeros([0], { dtype, device });
+  } else if (num === 1) {
+    return scalar(start, { dtype, device }).reshape([1]);
+  } else if (start === stop) {
+    return full([num], start, { dtype, device });
+  }
+  const delta = stop - start;
+  const denom = endpoint ? num - 1 : num;
+  const exp3 = AluExp.cast(
+    dtype,
+    AluExp.add(
+      AluExp.f32(start),
+      AluExp.mul(
+        AluExp.f32(delta / denom),
+        AluExp.cast("float32" /* Float32 */, AluVar.idx)
+      )
+    )
+  );
+  const st = ShapeTracker.fromShape([num]);
+  return new Array3(exp3, st, dtype, getBackend(device));
+}
+function aluCompare(a, b, op) {
+  switch (op) {
+    case "greater" /* Greater */:
+      return AluExp.mul(AluExp.cmpne(a, b), AluExp.cmplt(a, b).not());
+    case "less" /* Less */:
+      return AluExp.cmplt(a, b);
+    case "equal" /* Equal */:
+      return AluExp.cmpne(a, b).not();
+    case "not_equal" /* NotEqual */:
+      return AluExp.cmpne(a, b);
+    case "greater_equal" /* GreaterEqual */:
+      return AluExp.cmplt(a, b).not();
+    case "less_equal" /* LessEqual */:
+      return AluExp.add(AluExp.cmplt(a, b), AluExp.cmpne(a, b).not());
+  }
+}
+function generalBroadcast(a, b) {
+  const out = [];
+  let i = a.length - 1;
+  let j = b.length - 1;
+  for (; i >= 0 && j >= 0; i--, j--) {
+    const x = a[i];
+    const y = b[j];
+    if (x === y) {
+      out.push(x);
+    } else if (x === 1) {
+      out.push(y);
+    } else if (y === 1) {
+      out.push(x);
+    } else {
+      throw new TypeError(`Incompatible array broadcast shapes: ${a} vs ${b}`);
+    }
+  }
+  for (; i >= 0; i--) {
+    out.push(a[i]);
+  }
+  for (; j >= 0; j--) {
+    out.push(b[j]);
+  }
+  return out.reverse();
+}
+
+// src/frontend/jaxpr.ts
+var Var = class _Var {
+  static #nextId = 1;
+  // For debugging, since JavaScript has no id() function like Python.
+  id;
+  aval;
+  constructor(aval) {
+    this.id = _Var.#nextId++;
+    this.aval = aval;
+  }
+  toString() {
+    return `Var(${this.id}):${this.aval.strShort()}`;
+  }
+};
+var Lit = class {
+  dtype;
+  value;
+  aval;
+  constructor(dtype, value) {
+    this.dtype = dtype;
+    this.value = value;
+    this.aval = new ShapedArray([], dtype);
+  }
+};
+function atomIsLit(atom, literal) {
+  return atom instanceof Lit && (literal === void 0 || atom.value === literal);
+}
+var VarPrinter = class {
+  names = /* @__PURE__ */ new Map();
+  #next = "a";
+  // a, b, c, ..., z, aa, ab, ..., az, ba, bb, ...
+  #advance() {
+    const ret = this.#next;
+    let lastNonz = this.#next.length - 1;
+    while (lastNonz >= 0 && this.#next[lastNonz] === "z") {
+      lastNonz--;
+    }
+    if (lastNonz < 0) {
+      this.#next = "a".repeat(this.#next.length + 1);
+    } else {
+      let result = this.#next.slice(0, lastNonz);
+      result += String.fromCharCode(this.#next.charCodeAt(lastNonz) + 1);
+      result += "a".repeat(this.#next.length - 1 - lastNonz);
+      this.#next = result;
+    }
+    return ret;
+  }
+  name(v) {
+    if (this.names.has(v)) {
+      return this.names.get(v);
+    }
+    const name = this.#advance();
+    this.names.set(v, name);
+    return name;
+  }
+  nameType(v) {
+    return `${this.name(v)}:${v.aval.strShort()}`;
+  }
+};
+var JaxprEqn = class {
+  constructor(primitive, inputs, params, outBinders) {
+    this.primitive = primitive;
+    this.inputs = inputs;
+    this.params = params;
+    this.outBinders = outBinders;
+  }
+  pprint(usedVars, vp = new VarPrinter()) {
+    const lhs = PPrint.pp(
+      this.outBinders.map((v) => !usedVars || usedVars.has(v) ? vp.nameType(v) : "_").join(" ")
+    );
+    let rhs = PPrint.pp(this.primitive);
+    const paramsList = Object.entries(this.params).map(
+      ([k, v]) => PPrint.pp(`${k}=${v}`)
+    );
+    if (paramsList.length > 0) {
+      rhs = rhs.stack(PPrint.pp(" [ ")).stack(PPrint.prototype.concat(...paramsList)).stack(PPrint.pp(" ] "));
+    } else {
+      rhs = rhs.stack(PPrint.pp(" "));
+    }
+    rhs = rhs.stack(
+      PPrint.pp(
+        this.inputs.map((x) => x instanceof Var ? vp.name(x) : JSON.stringify(x.value)).join(" ")
+      )
+    );
+    return lhs.stack(PPrint.pp(" = ")).stack(rhs);
+  }
+  toString() {
+    return this.pprint().toString();
+  }
+};
+var Jaxpr2 = class _Jaxpr {
+  constructor(inBinders, eqns, outs) {
+    this.inBinders = inBinders;
+    this.eqns = eqns;
+    this.outs = outs;
+  }
+  #hash;
+  pprint() {
+    const vp = new VarPrinter();
+    const usedVars = new Set(
+      [...this.outs, ...this.eqns.flatMap((eqn) => eqn.inputs)].filter(
+        (x) => x instanceof Var
+      )
+    );
+    const inBinders = this.inBinders.map((v) => vp.nameType(v)).join(", ");
+    const eqns = PPrint.prototype.concat(
+      ...this.eqns.map((e2) => e2.pprint(usedVars, vp))
+    );
+    const outs = this.outs.map((x) => x instanceof Var ? vp.name(x) : x.value).join(", ");
+    return PPrint.pp(`{ lambda ${inBinders} .`).concat(
+      (this.eqns.length ? PPrint.pp("let ").stack(eqns).concat(PPrint.pp(`in ( ${outs} ) }`)) : PPrint.pp(`( ${outs} ) }`)).indent(2)
+    );
+  }
+  toString() {
+    return this.pprint().toString();
+  }
+  /**
+   * Gets a hash of this Jaxpr.
+   *
+   * Var identity is not considered in the hash, so two Jaxprs with the same
+   * order of assignments and operators but different variable IDs will resolve
+   * to the same hash (and toString representation).
+   */
+  getHash() {
+    if (this.#hash !== void 0) return this.#hash;
+    const hasher = new FpHash();
+    const varIds = /* @__PURE__ */ new Map();
+    const vi = (v) => {
+      if (varIds.has(v)) return varIds.get(v);
+      const id = varIds.size + 1;
+      varIds.set(v, FpHash.hash(id, v.aval.dtype, ...v.aval.shape));
+      return id;
+    };
+    hasher.update(this.inBinders.length, ...this.inBinders.map(vi));
+    hasher.update(
+      this.eqns.length,
+      ...this.eqns.flatMap((eqn) => [
+        eqn.primitive,
+        eqn.inputs.length,
+        ...eqn.inputs.map((x) => x instanceof Var ? vi(x) : x.value),
+        JSON.stringify(eqn.params),
+        eqn.outBinders.length,
+        ...eqn.outBinders.map(vi)
+      ])
+    );
+    hasher.update(
+      this.outs.length,
+      ...this.outs.map((x) => x instanceof Var ? vi(x) : x.value)
+    );
+    return this.#hash = hasher.value;
+  }
+  hash(state) {
+    state.update(this.getHash());
+  }
+  /**
+   * Produce a simplified Jaxpr with basic optimizations applied.
+   *  - Trim away unused variables.
+   *  - Fold away *1, *0, or +0 operations against literals.
+   */
+  simplify() {
+    const context = /* @__PURE__ */ new Map();
+    const newEqns = [];
+    for (const e2 of this.eqns) {
+      const inputs = e2.inputs.map(
+        (x) => x instanceof Var ? context.get(x) ?? x : x
+      );
+      const eqn = new JaxprEqn(e2.primitive, inputs, e2.params, e2.outBinders);
+      if (eqn.primitive === "add" /* Add */) {
+        const [a, b] = inputs;
+        const c = eqn.outBinders[0];
+        if (atomIsLit(a, 0)) {
+          context.set(c, b);
+        } else if (atomIsLit(b, 0)) {
+          context.set(c, a);
+        } else if (atomIsLit(a) && atomIsLit(b)) {
+          context.set(
+            c,
+            new Lit(
+              a.dtype,
+              a.dtype === "bool" /* Bool */ ? Math.min(a.value + b.value, 1) : a.value + b.value
+            )
+          );
+        } else {
+          newEqns.push(eqn);
+        }
+      } else if (eqn.primitive === "mul" /* Mul */) {
+        const [a, b] = inputs;
+        const c = eqn.outBinders[0];
+        if (atomIsLit(a, 1)) {
+          context.set(c, b);
+        } else if (atomIsLit(b, 1)) {
+          context.set(c, a);
+        } else if (atomIsLit(a) && atomIsLit(b)) {
+          context.set(c, new Lit(a.dtype, a.value * b.value));
+        } else {
+          newEqns.push(eqn);
+        }
+      } else if (eqn.primitive === "idiv" /* Idiv */) {
+        const [a, b] = inputs;
+        const c = eqn.outBinders[0];
+        if (atomIsLit(b, 1)) {
+          context.set(c, a);
+        }
+      } else {
+        newEqns.push(eqn);
+      }
+    }
+    const outs = this.outs.map(
+      (x) => x instanceof Var ? context.get(x) ?? x : x
+    );
+    const usedVars = new Set(outs.filter((x) => x instanceof Var));
+    const liveEqns = [];
+    for (let i = newEqns.length - 1; i >= 0; i--) {
+      const eqn = newEqns[i];
+      if (eqn.outBinders.some((v) => usedVars.has(v))) {
+        liveEqns.push(eqn);
+        for (const v of eqn.inputs) {
+          if (v instanceof Var) {
+            usedVars.add(v);
+          }
+        }
+      }
+    }
+    return new _Jaxpr(this.inBinders, liveEqns.reverse(), outs);
+  }
+  /** Flattens nested JitCall in a Jaxpr. Useful for handling jit-of-jit. */
+  flatten() {
+    if (!this.eqns.some((eqn) => eqn.primitive === "jit_call" /* JitCall */)) {
+      return this;
+    }
+    const newEqns = [];
+    const varMap = /* @__PURE__ */ new Map();
+    const varMapF = (x) => x instanceof Var ? varMap.get(x) ?? x : x;
+    for (const eqn of this.eqns) {
+      if (eqn.primitive === "jit_call" /* JitCall */) {
+        const jaxpr = eqn.params.jaxpr.flatten();
+        const translation = /* @__PURE__ */ new Map();
+        const translationF = (x) => x instanceof Var ? translation.get(x) : x;
+        for (const [v, x] of zip(jaxpr.inBinders, eqn.inputs)) {
+          translation.set(v, varMapF(x));
+        }
+        for (const ieqn of jaxpr.eqns) {
+          const inputs = ieqn.inputs.map(translationF);
+          const outBinders = [];
+          for (const v of ieqn.outBinders) {
+            const u = new Var(v.aval);
+            outBinders.push(u);
+            translation.set(v, u);
+          }
+          newEqns.push(
+            new JaxprEqn(ieqn.primitive, inputs, ieqn.params, outBinders)
+          );
+        }
+        for (const [v, x] of zip(eqn.outBinders, jaxpr.outs)) {
+          varMap.set(v, translationF(x));
+        }
+      } else {
+        if (eqn.inputs.some((x) => x instanceof Var && varMap.has(x))) {
+          newEqns.push(
+            new JaxprEqn(
+              eqn.primitive,
+              eqn.inputs.map(varMapF),
+              eqn.params,
+              eqn.outBinders
+            )
+          );
+        } else {
+          newEqns.push(eqn);
+        }
+      }
+    }
+    const newOuts = this.outs.map(varMapF);
+    return new _Jaxpr(this.inBinders, newEqns, newOuts);
+  }
+};
+var JaxprType = class {
+  constructor(inTypes, outTypes) {
+    this.inTypes = inTypes;
+    this.outTypes = outTypes;
+  }
+  toString() {
+    const inTypes = this.inTypes.map((aval) => aval.strShort()).join(", ");
+    const outTypes = this.outTypes.map((aval) => aval.strShort()).join(", ");
+    return `(${inTypes}) -> (${outTypes})`;
+  }
+};
+function typecheckJaxpr(jaxpr) {
+  const env = /* @__PURE__ */ new Set();
+  for (const v of jaxpr.inBinders) {
+    if (env.has(v)) {
+      throw new TypeError(`Duplicate variable binding: ${v}`);
+    }
+    env.add(v);
+  }
+  for (const eqn of jaxpr.eqns) {
+    const inTypes2 = eqn.inputs.map((x) => typecheckAtom(env, x));
+    const outTypes2 = abstractEvalRules[eqn.primitive](inTypes2, eqn.params);
+    for (const [outBinder, outType] of zip(eqn.outBinders, outTypes2)) {
+      if (!outType.equals(outBinder.aval)) {
+        throw new TypeError(
+          `Output binder type mismatch in ${eqn.primitive}: ${outBinder} vs ${outType}`
+        );
+      }
+      if (env.has(outBinder)) {
+        throw new TypeError(`Duplicate variable binding: ${outBinder}`);
+      }
+      env.add(outBinder);
+    }
+  }
+  const inTypes = jaxpr.inBinders.map((v) => v.aval);
+  const outTypes = jaxpr.outs.map((x) => typecheckAtom(env, x));
+  return new JaxprType(inTypes, outTypes);
+}
+function typecheckAtom(env, x) {
+  if (x instanceof Var) {
+    if (!env.has(x)) {
+      throw new Error(`Unknown variable: ${x}`);
+    }
+    return x.aval;
+  } else if (x instanceof Lit) {
+    return x.aval;
+  } else {
+    throw new TypeError(`Invalid atom type: ${x}`);
+  }
+}
+function evalJaxpr(jaxpr, args) {
+  const env = /* @__PURE__ */ new Map();
+  const usageCount = /* @__PURE__ */ new Map();
+  for (const x of jaxpr.eqns.flatMap((eqn) => eqn.inputs).concat(jaxpr.outs)) {
+    if (x instanceof Var) usageCount.set(x, (usageCount.get(x) ?? 0) + 1);
+  }
+  const remainingRefs = /* @__PURE__ */ new Map();
+  const read = (x) => {
+    if (x instanceof Var) {
+      remainingRefs.set(x, (remainingRefs.get(x) ?? 0) - 1);
+      return env.get(x);
+    } else {
+      return scalar(x.value, { dtype: x.dtype });
+    }
+  };
+  const write = (v, val) => {
+    if (env.has(v)) throw new Error(`Variable already bound: ${v}`);
+    let refCount = usageCount.get(v) ?? 0;
+    if (refCount) {
+      env.set(v, val);
+      remainingRefs.set(v, refCount);
+      while (refCount-- > 1) val.ref;
+    } else {
+      val.dispose();
+    }
+  };
+  try {
+    for (const [v, arg] of zip(jaxpr.inBinders, args)) write(v, arg);
+    for (const eqn of jaxpr.eqns) {
+      const inVals = eqn.inputs.map(read);
+      const outVals = bind(eqn.primitive, inVals, eqn.params);
+      for (const [v, val] of zip(eqn.outBinders, outVals)) write(v, val);
+    }
+    return jaxpr.outs.map(read);
+  } catch (error) {
+    for (let [v, refCount] of remainingRefs.entries()) {
+      if (refCount > 0) {
+        const tracer = env.get(v);
+        while (refCount--) tracer.dispose();
+      }
+    }
+    throw error;
+  }
+}
+function jaxprAsFun(jaxpr) {
+  return (...args) => evalJaxpr(jaxpr, args);
+}
+var JaxprTracer = class extends Tracer {
+  constructor(trace, aval) {
+    super(trace);
+    this.aval = aval;
+  }
+  toString() {
+    return `JaxprTracer(${this.aval.strShort()})`;
+  }
+  // JaxprTracer does not hold any resources, no need to be reference counted.
+  get ref() {
+    return this;
+  }
+  dispose() {
+  }
+};
+var JaxprTrace = class extends Trace {
+  /** Register a Jaxpr argument with a given shape and return the tracer. */
+  newArg(aval) {
+    aval = ShapedArray.fromAval(aval);
+    const tracer = this.builder.newTracer(this, aval);
+    this.builder.addVar(tracer);
+    return tracer;
+  }
+  /** Register a constant / literal in this Jaxpr. */
+  getOrMakeConstTracer(val) {
+    let tracer = this.builder.constTracers.get(val);
+    if (tracer === void 0) {
+      tracer = this.builder.newTracer(this, ShapedArray.fromAval(getAval(val)));
+      this.builder.addConst(
+        tracer,
+        val instanceof Tracer ? val.ref : scalar(val)
+      );
+    }
+    return tracer;
+  }
+  pure = this.getOrMakeConstTracer;
+  lift = this.getOrMakeConstTracer;
+  processPrimitive(primitive, tracers, params) {
+    const avalsIn = tracers.map((t) => t.aval);
+    const avalsOut = abstractEvalRules[primitive](avalsIn, params);
+    const outTracers = avalsOut.map(
+      (aval) => this.builder.newTracer(this, aval)
+    );
+    this.builder.addEqn(
+      new JaxprEqn(
+        primitive,
+        tracers.map((t) => this.builder.getVar(t)),
+        params,
+        outTracers.map((t) => this.builder.addVar(t))
+      )
+    );
+    return outTracers;
+  }
+  get builder() {
+    return this.main.globalData;
+  }
+};
+var JaxprBuilder = class {
+  eqns = [];
+  tracerToVar = /* @__PURE__ */ new Map();
+  constTracers = /* @__PURE__ */ new Map();
+  // already-seen value -> tracer
+  constVals = /* @__PURE__ */ new Map();
+  // var -> const value
+  tracers = [];
+  newTracer(trace, aval) {
+    const tracer = new JaxprTracer(trace, aval);
+    this.tracers.push(tracer);
+    return tracer;
+  }
+  addEqn(eqn) {
+    this.eqns.push(eqn);
+  }
+  addVar(tracer) {
+    if (this.tracerToVar.has(tracer)) {
+      throw new Error(`Tracer was added as variable twice: ${tracer}`);
+    }
+    const v = new Var(tracer.aval);
+    this.tracerToVar.set(tracer, v);
+    return v;
+  }
+  getVar(tracer) {
+    const v = this.tracerToVar.get(tracer);
+    if (v === void 0) {
+      throw new Error(`Could not find variable for tracer: ${tracer}`);
+    }
+    return v;
+  }
+  addConst(tracer, val) {
+    const v = this.addVar(tracer);
+    this.constTracers.set(val, tracer);
+    this.constVals.set(v, val);
+    return v;
+  }
+  build(inTracers, outTracers) {
+    let [constVars, consts] = unzip2(this.constVals.entries());
+    const t2v = this.getVar.bind(this);
+    const inBinders = [...constVars, ...inTracers.map(t2v)];
+    const outVars = outTracers.map(t2v);
+    let jaxpr = new Jaxpr2(inBinders, this.eqns, outVars);
+    typecheckJaxpr(jaxpr);
+    [jaxpr, consts] = _inlineLiterals(jaxpr, consts);
+    return { jaxpr, consts };
+  }
+};
+function _inlineLiterals(jaxpr, consts) {
+  const literals = /* @__PURE__ */ new Map();
+  const constBinders = [];
+  const newConsts = [];
+  for (let i = 0; i < consts.length; i++) {
+    if (ndim(consts[i]) === 0 && consts[i] instanceof Array3) {
+      const ar = consts[i];
+      literals.set(jaxpr.inBinders[i], new Lit(ar.dtype, ar.dataSync()[0]));
+    } else {
+      constBinders.push(jaxpr.inBinders[i]);
+      newConsts.push(consts[i]);
+    }
+  }
+  const newEqns = jaxpr.eqns.map(
+    (eqn) => new JaxprEqn(
+      eqn.primitive,
+      eqn.inputs.map((x) => literals.get(x) ?? x),
+      eqn.params,
+      eqn.outBinders
+    )
+  );
+  const newOuts = jaxpr.outs.map((x) => literals.get(x) ?? x);
+  const newJaxpr = new Jaxpr2(
+    [...constBinders, ...jaxpr.inBinders.slice(consts.length)],
+    newEqns,
+    newOuts
+  );
+  typecheckJaxpr(newJaxpr);
+  return [newJaxpr, newConsts];
+}
+function binopAbstractEval([x, y]) {
+  if (!(x instanceof ShapedArray) || !(y instanceof ShapedArray)) {
+    throw new TypeError("binopAbstractEval expects ShapedArray inputs");
+  }
+  if (x.dtype !== y.dtype) {
+    throw new TypeError(`Mismatched dtypes: ${x.dtype} vs ${y.dtype}`);
+  }
+  return [new ShapedArray(generalBroadcast(x.shape, y.shape), x.dtype)];
+}
+function compareAbstractEval([x, y]) {
+  if (!(x instanceof ShapedArray) || !(y instanceof ShapedArray)) {
+    throw new TypeError("compareAbstractEval expects ShapedArray inputs");
+  }
+  if (x.dtype !== y.dtype) {
+    throw new TypeError(`Mismatched dtypes: ${x.dtype} vs ${y.dtype}`);
+  }
+  return [new ShapedArray(generalBroadcast(x.shape, y.shape), "bool" /* Bool */)];
+}
+function vectorizedUnopAbstractEval([x]) {
+  return [ShapedArray.fromAval(x)];
+}
+var abstractEvalRules = {
+  ["add" /* Add */]: binopAbstractEval,
+  ["mul" /* Mul */]: binopAbstractEval,
+  ["idiv" /* Idiv */]: binopAbstractEval,
+  ["neg" /* Neg */]: vectorizedUnopAbstractEval,
+  ["reciprocal" /* Reciprocal */]: vectorizedUnopAbstractEval,
+  ["sin" /* Sin */]: vectorizedUnopAbstractEval,
+  ["cos" /* Cos */]: vectorizedUnopAbstractEval,
+  ["exp" /* Exp */]: vectorizedUnopAbstractEval,
+  ["log" /* Log */]: vectorizedUnopAbstractEval,
+  ["min" /* Min */]: binopAbstractEval,
+  ["max" /* Max */]: binopAbstractEval,
+  ["reduce_sum" /* ReduceSum */]([x], { axis }) {
+    const axisSet = new Set(axis);
+    const newShape = x.shape.filter((_, i) => !axisSet.has(i));
+    return [new ShapedArray(newShape, x.dtype)];
+  },
+  ["compare" /* Compare */]: compareAbstractEval,
+  ["where" /* Where */]([cond, x, y]) {
+    if (cond.dtype !== "bool" /* Bool */)
+      throw new TypeError(`Condition must be boolean, got ${cond.dtype}`);
+    if (x.dtype !== y.dtype)
+      throw new TypeError(`Mismatched dtypes: ${x.dtype} vs ${y.dtype}`);
+    const shape2 = generalBroadcast(
+      cond.shape,
+      generalBroadcast(x.shape, y.shape)
+    );
+    return [new ShapedArray(shape2, x.dtype)];
+  },
+  ["transpose" /* Transpose */]([x], { perm }) {
+    return [
+      new ShapedArray(
+        perm.map((i) => x.shape[i]),
+        x.dtype
+      )
+    ];
+  },
+  ["broadcast" /* Broadcast */]([x], { shape: shape2 }) {
+    return [new ShapedArray(shape2, x.dtype)];
+  },
+  ["reshape" /* Reshape */]([x], { shape: shape2 }) {
+    return [new ShapedArray(shape2, x.dtype)];
+  },
+  ["flip" /* Flip */]([x], _) {
+    return [new ShapedArray(x.shape, x.dtype)];
+  },
+  ["jit_call" /* JitCall */](args, { jaxpr }) {
+    const { inTypes, outTypes } = typecheckJaxpr(jaxpr);
+    if (args.length !== inTypes.length) {
+      throw new TypeError(
+        `jit_call expected ${inTypes.length} arguments, got ${args.length}`
+      );
+    }
+    for (let i = 0; i < inTypes.length; i++) {
+      if (!args[i].equals(inTypes[i])) {
+        throw new TypeError(
+          `jit_call argument ${i} has type ${args[i]}, expected ${inTypes[i]}`
+        );
+      }
+    }
+    return outTypes;
+  }
+};
+function makeJaxpr(f) {
+  return (...argsIn) => {
+    var _stack = [];
+    try {
+      const [avalsIn, inTree] = flatten(argsIn);
+      const [fFlat, outTree] = flattenFun(f, inTree);
+      const builder = new JaxprBuilder();
+      const main = __using(_stack, newMain(JaxprTrace, builder));
+      const _dynamic = __using(_stack, newDynamic(main));
+      const trace = new JaxprTrace(main);
+      const tracersIn = avalsIn.map(
+        (aval) => trace.newArg(typeof aval === "object" ? aval : pureArray(aval))
+      );
+      const outs = fFlat(...tracersIn);
+      const tracersOut = outs.map(
+        (out) => fullRaise(trace, out)
+      );
+      const { jaxpr, consts } = builder.build(tracersIn, tracersOut);
+      if (outTree.value === void 0) {
+        throw new Error("outTree was not set in makeJaxpr");
+      }
+      return { jaxpr: jaxpr.simplify(), consts, treedef: outTree.value };
+    } catch (_) {
+      var _error = _, _hasError = true;
+    } finally {
+      __callDispose(_stack, _error, _hasError);
+    }
+  };
+}
+function jit(f) {
+  const cache = /* @__PURE__ */ new Map();
+  return (...args) => {
+    const [argsFlat, inTree] = flatten(args);
+    const avalsInFlat = argsFlat.map((x) => ShapedArray.fromAval(getAval(x)));
+    const avalsIn = unflatten(inTree, avalsInFlat);
+    const cacheKey = JSON.stringify(avalsIn);
+    const {
+      jaxpr,
+      consts,
+      treedef: outTree
+    } = runWithCache(
+      cache,
+      cacheKey,
+      () => makeJaxpr(f)(...avalsIn)
+    );
+    const outs = bind(
+      "jit_call" /* JitCall */,
+      [...consts.map((c) => c.ref), ...argsFlat],
+      {
+        jaxpr,
+        numConsts: consts.length
+      }
+    );
+    return unflatten(outTree, outs);
+  };
+}
+
+// src/frontend/jvp.ts
+var JVPTracer = class extends Tracer {
+  constructor(trace, primal, tangent) {
+    super(trace);
+    this.primal = primal;
+    this.tangent = tangent;
+  }
+  get aval() {
+    return this.primal.aval;
+  }
+  toString() {
+    return `JVPTracer(${this.primal.toString()}, ${this.tangent.toString()})`;
+  }
+  get ref() {
+    this.primal.ref, this.tangent.ref;
+    return this;
+  }
+  dispose() {
+    this.primal.dispose();
+    this.tangent.dispose();
+  }
+};
+var JVPTrace = class extends Trace {
+  pure(val) {
+    return this.lift(pureArray(val));
+  }
+  lift(val) {
+    return new JVPTracer(this, val, zerosLike(val));
+  }
+  processPrimitive(primitive, tracers, params) {
+    const [primalsIn, tangentsIn] = unzip2(
+      tracers.map((x) => [x.primal, x.tangent])
+    );
+    const jvpRule = jvpRules[primitive];
+    if (jvpRule === void 0) {
+      throw new Error(`No JVP rule for: ${primitive}`);
+    }
+    const [primalsOut, tangentsOut] = jvpRule(primalsIn, tangentsIn, params);
+    return zip(primalsOut, tangentsOut).map(
+      ([x, t]) => new JVPTracer(this, x, t)
+    );
+  }
+};
+var jvpRules = {
+  ["add" /* Add */]([x, y], [dx, dy]) {
+    return [[x.add(y)], [dx.add(dy)]];
+  },
+  ["mul" /* Mul */]([x, y], [dx, dy]) {
+    return [[x.ref.mul(y.ref)], [x.mul(dy).add(dx.mul(y))]];
+  },
+  ["idiv" /* Idiv */]([x, y], [dx, dy]) {
+    dx.dispose(), dy.dispose();
+    const z = idiv(x, y);
+    const dz = zerosLike(z);
+    return [[z], [dz]];
+  },
+  ["neg" /* Neg */]([x], [dx]) {
+    return [[x.neg()], [dx.neg()]];
+  },
+  ["reciprocal" /* Reciprocal */]([x], [dx]) {
+    const xRecip = reciprocal(x.ref);
+    return [[xRecip.ref], [neg(xRecip.ref.mul(xRecip)).mul(dx)]];
+  },
+  ["sin" /* Sin */]([x], [dx]) {
+    return [[sin(x.ref)], [cos(x).mul(dx)]];
+  },
+  ["cos" /* Cos */]([x], [dx]) {
+    return [[cos(x.ref)], [neg(sin(x)).mul(dx)]];
+  },
+  ["exp" /* Exp */]([x], [dx]) {
+    const z = exp(x);
+    return [[z.ref], [z.mul(dx)]];
+  },
+  ["log" /* Log */]([x], [dx]) {
+    return [[log(x.ref)], [reciprocal(x).mul(dx)]];
+  },
+  ["min" /* Min */]([x, y], [dx, dy]) {
+    return [[min(x.ref, y.ref)], [where(less(y, x), dy, dx)]];
+  },
+  ["max" /* Max */]([x, y], [dx, dy]) {
+    return [[max(x.ref, y.ref)], [where(less(y, x), dx, dy)]];
+  },
+  ["reduce_sum" /* ReduceSum */]([x], [dx], { axis }) {
+    return [[reduceSum(x, axis)], [reduceSum(dx, axis)]];
+  },
+  ["compare" /* Compare */]([x, y], tangents, { op }) {
+    for (const t of tangents) t.dispose();
+    const primal = compare(x, y, op);
+    return [[primal], [zerosLike(primal)]];
+  },
+  ["where" /* Where */]([cond, x, y], [dcond, dx, dy]) {
+    dcond.dispose();
+    return [[where(cond.ref, x, y)], [where(cond, dx, dy)]];
+  },
+  ["transpose" /* Transpose */]([x], [dx], { perm }) {
+    return [[transpose(x, perm)], [transpose(dx, perm)]];
+  },
+  ["broadcast" /* Broadcast */]([x], [dx], { shape: shape2, axis }) {
+    return [[broadcast(x, shape2, axis)], [broadcast(dx, shape2, axis)]];
+  },
+  ["reshape" /* Reshape */]([x], [dx], { shape: shape2 }) {
+    return [[reshape(x, shape2)], [reshape(dx, shape2)]];
+  },
+  ["flip" /* Flip */]([x], [dx], { axis }) {
+    return [[flip(x, axis)], [flip(dx, axis)]];
+  },
+  ["jit_call" /* JitCall */](primals, tangents, { jaxpr }) {
+    const { newJaxpr, newConsts } = jvpJaxpr(jaxpr);
+    const outs = bind(
+      "jit_call" /* JitCall */,
+      [...newConsts, ...primals, ...tangents],
+      {
+        jaxpr: newJaxpr,
+        numConsts: newConsts.length
+      }
+    );
+    const n = outs.length / 2;
+    if (!Number.isInteger(n))
+      throw new Error("internal: JVP Jaxpr output length is not even");
+    const [primalsOut, tangentsOut] = [outs.slice(0, n), outs.slice(n)];
+    return [primalsOut, tangentsOut];
+  }
+};
+var jvpJaxprCache = /* @__PURE__ */ new Map();
+function jvpJaxpr(jaxpr) {
+  if (jvpJaxprCache.has(jaxpr)) {
+    return jvpJaxprCache.get(jaxpr);
+  }
+  const inAvals = jaxpr.inBinders.map((v) => v.aval);
+  const { jaxpr: newJaxpr, consts: newConsts } = makeJaxpr(
+    (primals, tangents) => jvpFlat(jaxprAsFun(jaxpr), primals, tangents)
+  )(inAvals, inAvals);
+  const result = { newJaxpr, newConsts };
+  jvpJaxprCache.set(jaxpr, result);
+  return result;
+}
+function jvpFlat(f, primals, tangents) {
+  var _stack = [];
+  try {
+    const main = __using(_stack, newMain(JVPTrace));
+    const trace = new JVPTrace(main);
+    const tracersIn = zip(primals, tangents).map(
+      ([x, t]) => new JVPTracer(trace, pureArray(x), pureArray(t))
+    );
+    const outs = f(...tracersIn);
+    const tracersOut = outs.map((out) => fullRaise(trace, out));
+    return unzip2(tracersOut.map((t) => [t.primal, t.tangent]));
+  } catch (_) {
+    var _error = _, _hasError = true;
+  } finally {
+    __callDispose(_stack, _error, _hasError);
+  }
+}
+function jvp(f, primals, tangents) {
+  const [primalsFlat, inTree] = flatten(primals);
+  const [tangentsFlat, inTree2] = flatten(tangents);
+  if (!inTree.equals(inTree2)) {
+    throw new TreeMismatchError("jvp", inTree, inTree2);
+  }
+  const [flatFun, outTree] = flattenFun(f, inTree);
+  const [primalsOutFlat, tangentsOutFlat] = jvpFlat(
+    flatFun,
+    primalsFlat,
+    tangentsFlat
+  );
+  if (outTree.value === void 0) {
+    throw new Error("outTree was not set in jvp");
+  }
+  const primalsOut = unflatten(outTree.value, primalsOutFlat);
+  const tangentsOut = unflatten(outTree.value, tangentsOutFlat);
+  return [primalsOut, tangentsOut];
+}
+
+// src/frontend/vmap.ts
+function mappedAval(batchDim, aval) {
+  const shape2 = [...aval.shape];
+  shape2.splice(batchDim, 1);
+  return new ShapedArray(shape2, aval.dtype);
+}
+function moveaxis(x, src, dst) {
+  const t = pureArray(x);
+  const perm = range(t.shape.length);
+  perm.splice(src, 1);
+  perm.splice(dst, 0, src);
+  return transpose(t, perm);
+}
+function moveBatchAxis(axisSize, src, dst, x) {
+  if (src === null) {
+    const targetShape = [...x.shape];
+    targetShape.splice(dst, 0, axisSize);
+    return broadcast(x, targetShape, [dst]);
+  } else if (src === dst) {
+    return x;
+  } else {
+    return moveaxis(x, src, dst);
+  }
+}
+var BatchTracer = class extends Tracer {
+  constructor(trace, val, batchDim) {
+    super(trace);
+    this.val = val;
+    this.batchDim = batchDim;
+  }
+  get aval() {
+    if (this.batchDim === null) {
+      return this.val.aval;
+    } else {
+      return mappedAval(this.batchDim, this.val.aval);
+    }
+  }
+  toString() {
+    return `BatchTracer(${this.val.toString()}, ${this.batchDim})`;
+  }
+  get ref() {
+    this.val.ref;
+    return this;
+  }
+  dispose() {
+    this.val.dispose();
+  }
+  fullLower() {
+    if (this.batchDim === null) {
+      return this.val.fullLower();
+    } else {
+      return this;
+    }
+  }
+};
+var BatchTrace = class extends Trace {
+  pure(val) {
+    return this.lift(pureArray(val));
+  }
+  lift(val) {
+    return new BatchTracer(this, val, null);
+  }
+  processPrimitive(primitive, tracers, params) {
+    const [valsIn, bdimsIn] = unzip2(tracers.map((t) => [t.val, t.batchDim]));
+    const vmapRule = vmapRules[primitive];
+    if (vmapRule === void 0) {
+      throw new Error(`No vmap rule for: ${primitive}`);
+    }
+    const [valOuts, bdimOuts] = vmapRule(
+      this.axisSize,
+      valsIn,
+      bdimsIn,
+      params
+    );
+    return zip(valOuts, bdimOuts).map(
+      ([x, bd]) => new BatchTracer(this, x, bd)
+    );
+  }
+  get axisSize() {
+    return this.main.globalData;
+  }
+};
+function handleScalarBroadcasting(nd, x, d) {
+  if (d === null || nd === ndim(x)) {
+    return x;
+  } else {
+    const axis = range(ndim(x), nd);
+    const shape2 = [...x.shape, ...axis.map(() => 1)];
+    return broadcast(x, shape2, axis);
+  }
+}
+function broadcastBatcher(op) {
+  return (axisSize, args, dims) => {
+    if (args.length === 0) {
+      throw new Error("Empty list in broadcastBatcher");
+    }
+    const idx = dims.findIndex((d) => d !== null);
+    if (idx === -1) {
+      return [[op(...args)], [null]];
+    }
+    if (
+      // If only agreeing batch dims, or scalars, just call the primitive.
+      zip(args, dims).every(
+        ([x, d]) => ndim(x) === 0 || deepEqual(x.shape, args[idx].shape) && d === dims[idx]
+      )
+    ) {
+      return [[op(...args)], [dims[idx]]];
+    }
+    args = args.map(
+      (x, i) => ndim(x) > 0 ? moveBatchAxis(axisSize, dims[i], 0, x) : x
+    );
+    const nd = Math.max(...args.map(ndim));
+    args = args.map((x, i) => handleScalarBroadcasting(nd, x, dims[i]));
+    return [[op(...args)], [0]];
+  };
+}
+function vectorizedUnopBatchingRule(op) {
+  return (axisSize, [x], [xBdim]) => {
+    return [[op(x)], [xBdim]];
+  };
+}
+var vmapRules = {
+  ["add" /* Add */]: broadcastBatcher(add),
+  ["mul" /* Mul */]: broadcastBatcher(mul),
+  ["idiv" /* Idiv */]: broadcastBatcher(idiv),
+  ["neg" /* Neg */]: vectorizedUnopBatchingRule(neg),
+  ["reciprocal" /* Reciprocal */]: vectorizedUnopBatchingRule(reciprocal),
+  ["sin" /* Sin */]: vectorizedUnopBatchingRule(sin),
+  ["cos" /* Cos */]: vectorizedUnopBatchingRule(cos),
+  ["exp" /* Exp */]: vectorizedUnopBatchingRule(exp),
+  ["log" /* Log */]: vectorizedUnopBatchingRule(log),
+  ["min" /* Min */]: broadcastBatcher(min),
+  ["max" /* Max */]: broadcastBatcher(max),
+  ["reduce_sum" /* ReduceSum */](axisSize, [x], [xBdim], { axis }) {
+    if (xBdim === null) {
+      return [[reduceSum(x, axis)], [null]];
+    }
+    const newAxis = axis.map((ax) => ax + (xBdim <= ax ? 1 : 0));
+    const outBdim = xBdim - axis.filter((ax) => ax < xBdim).length;
+    return [[reduceSum(x, newAxis)], [outBdim]];
+  },
+  ["compare" /* Compare */](axisSize, args, dims, { op }) {
+    return broadcastBatcher((x, y) => compare(x, y, op))(axisSize, args, dims);
+  },
+  // TODO: where, transpose, broadcast, reshape, flip
+  ["jit_call" /* JitCall */](axisSize, args, dims, { jaxpr }) {
+    const { newJaxpr, newConsts } = vmapJaxpr(jaxpr, axisSize, dims);
+    const outs = bind("jit_call" /* JitCall */, [...newConsts, ...args], {
+      jaxpr: newJaxpr,
+      numConsts: newConsts.length
+    });
+    return [outs, rep(outs.length, 0)];
+  }
+};
+var vmapJaxprCache = /* @__PURE__ */ new Map();
+function vmapJaxpr(jaxpr, axisSize, dims) {
+  const cacheKey = JSON.stringify([axisSize, dims]);
+  const prevResult = vmapJaxprCache.get(jaxpr)?.get(cacheKey);
+  if (prevResult) return prevResult;
+  const inAvals = jaxpr.inBinders.map((v, i) => {
+    if (dims[i] === null) return v.aval;
+    const shape2 = [...v.aval.shape];
+    shape2.splice(dims[i], 0, axisSize);
+    return new ShapedArray(shape2, v.aval.dtype);
+  });
+  const { jaxpr: newJaxpr, consts: newConsts } = makeJaxpr(
+    (args) => vmapFlat(jaxprAsFun(jaxpr), dims, args)
+  )(inAvals);
+  const result = { newJaxpr, newConsts };
+  if (!vmapJaxprCache.has(jaxpr)) vmapJaxprCache.set(jaxpr, /* @__PURE__ */ new Map());
+  vmapJaxprCache.get(jaxpr).set(cacheKey, result);
+  return result;
+}
+function vmapFlat(f, inAxes, args) {
+  let axisSize = void 0;
+  for (let i = 0; i < args.length; i++) {
+    if (inAxes[i] !== null) {
+      const arg = args[i];
+      if (!(arg instanceof Tracer)) {
+        throw new TypeError("vmap requires Tracer argument for mapped axes");
+      }
+      const size2 = arg.shape[inAxes[i]];
+      if (axisSize === void 0) {
+        axisSize = size2;
+      } else if (axisSize !== size2) {
+        throw new TypeError(
+          "vmap requires all mapped axes to have the same size"
+        );
+      }
+    }
+  }
+  if (axisSize === void 0) {
+    throw new TypeError("vmap requires at least one mapped axis");
+  }
+  let valsOut, bdimsOut;
+  {
+    var _stack = [];
+    try {
+      const main = __using(_stack, newMain(BatchTrace, axisSize));
+      const trace = new BatchTrace(main);
+      const tracersIn = args.map(
+        (x, i) => inAxes[i] === null ? pureArray(x) : new BatchTracer(trace, pureArray(x), inAxes[i])
+      );
+      const outs = f(...tracersIn);
+      const tracersOut = outs.map((out) => fullRaise(trace, out));
+      [valsOut, bdimsOut] = unzip2(tracersOut.map((t) => [t.val, t.batchDim]));
+    } catch (_) {
+      var _error = _, _hasError = true;
+    } finally {
+      __callDispose(_stack, _error, _hasError);
+    }
+  }
+  return zip(valsOut, bdimsOut).map(
+    ([valOut, bdim]) => moveBatchAxis(axisSize, bdim, 0, valOut)
+  );
+}
+function vmap(f, inAxes = 0) {
+  return (...args) => {
+    const [argsFlat, inTree] = flatten(args);
+    let inAxesFlat;
+    if (typeof inAxes === "number") {
+      inAxesFlat = rep(argsFlat.length, inAxes);
+    } else {
+      let inTree2;
+      [inAxesFlat, inTree2] = flatten(inAxes);
+      if (!inTree.equals(inTree2)) {
+        throw new TreeMismatchError("vmap", inTree, inTree2);
+      }
+    }
+    const [fFlat, outTree] = flattenFun(f, inTree);
+    const outsFlat = vmapFlat(fFlat, inAxesFlat, argsFlat);
+    if (outTree.value === void 0) {
+      throw new Error("outTree was not set in vmap");
+    }
+    return unflatten(outTree.value, outsFlat);
+  };
+}
+function jacfwd(f) {
+  return function jacobianForward(x) {
+    if (x.shape.length !== 1) {
+      throw new TypeError("jacfwd only supports 1D inputs");
+    }
+    const [size2] = x.shape;
+    const pushfwd = (v) => jvp(f, [x], [v])[1];
+    return vmap(pushfwd, [0])(eye(size2, void 0, { dtype: x.dtype }));
+  };
+}
+
+// src/frontend/linearize.ts
+var PartialVal = class _PartialVal {
+  constructor(val, aval) {
+    this.val = val;
+    this.aval = aval;
+  }
+  static known(val) {
+    return new _PartialVal(val, ShapedArray.fromAval(val.aval));
+  }
+  static unknown(aval) {
+    return new _PartialVal(null, ShapedArray.fromAval(aval));
+  }
+  get isKnown() {
+    return this.val !== null;
+  }
+  toString() {
+    return this.val ? this.val.toString() : this.aval.strShort();
+  }
+};
+function partialEvalFlat(f, pvalsIn) {
+  const main = newMain(PartialEvalTrace);
+  const trace = new PartialEvalTrace(main);
+  const tracersIn = pvalsIn.map((pval) => trace.newArg(pval));
+  const unknownTracersIn = tracersIn.filter((t) => !t.pval.isKnown).map((t) => t.ref);
+  const outs = f(...tracersIn);
+  const tracersOut = outs.map(
+    (out) => fullRaise(trace, out)
+  );
+  const pvalsOut = tracersOut.map((t) => t.pval);
+  const unknownTracersOut = tracersOut.filter((t) => !t.pval.isKnown);
+  const { jaxpr, consts } = partialEvalGraphToJaxpr(
+    unknownTracersIn,
+    unknownTracersOut
+  );
+  return { jaxpr, pvalsOut, consts };
+}
+function linearizeFlatUtil(f, primalsIn) {
+  const pvalsIn = [
+    ...primalsIn.map(PartialVal.known),
+    ...primalsIn.map((t) => PartialVal.unknown(t.aval))
+  ];
+  const fJvp = (...x) => {
+    const k = x.length / 2;
+    const [primalsOut2, tangentsOut] = jvp(f, x.slice(0, k), x.slice(k, 2 * k));
+    return [...primalsOut2, ...tangentsOut];
+  };
+  const { jaxpr, pvalsOut, consts } = partialEvalFlat(fJvp, pvalsIn);
+  const primalPvals = pvalsOut.slice(0, pvalsOut.length / 2);
+  if (!primalPvals.every((pval) => pval.isKnown)) {
+    throw new TypeError(
+      "Not all primal values are known after partial evaluation"
+    );
+  }
+  const primalsOut = primalPvals.map((pval) => pval.val);
+  return { primalsOut, jaxpr, consts };
+}
+function linearizeFlat(f, primalsIn) {
+  const { primalsOut, jaxpr, consts } = linearizeFlatUtil(f, primalsIn);
+  const fLin = (...tangents) => evalJaxpr(jaxpr, [...consts.map((c) => c.ref), ...tangents]);
+  return [primalsOut, fLin];
+}
+function linearize(f, ...primalsIn) {
+  const [primalsInFlat, inTree] = flatten(primalsIn);
+  const [fFlat, outTree] = flattenFun(f, inTree);
+  const [primalsOutFlat, fLinFlat] = linearizeFlat(
+    fFlat,
+    primalsInFlat.map(pureArray)
+  );
+  if (outTree.value === void 0) {
+    throw new Error("outTree was not set in linearize");
+  }
+  const primalsOut = unflatten(outTree.value, primalsOutFlat);
+  const fLin = (...tangentsIn) => {
+    const [tangentsInFlat, inTree2] = flatten(tangentsIn);
+    if (!inTree.equals(inTree2)) {
+      throw new TreeMismatchError("linearize", inTree, inTree2);
+    }
+    const tangentsOutFlat = fLinFlat(...tangentsInFlat.map(pureArray));
+    return unflatten(outTree.value, tangentsOutFlat);
+  };
+  return [primalsOut, fLin];
+}
+var PartialEvalTracer = class extends Tracer {
+  // PartialEvalTracer reference count, used to free references.
+  // Note: Either pval is known and recipe is null, or pval is unknown and
+  // recipe describes how to compute the value.
+  constructor(trace, pval, recipe) {
+    super(trace);
+    this.pval = pval;
+    this.recipe = recipe;
+    this.#rc = 1;
+  }
+  #rc;
+  get aval() {
+    return this.pval.aval;
+  }
+  toString() {
+    if (!this.recipe) {
+      return `PartialEvalTracer(${this.pval.toString()})`;
+    } else {
+      return `PartialEvalTracer<${this.recipe.type}>(${this.pval.toString()})`;
+    }
+  }
+  get ref() {
+    if (this.#rc <= 0) {
+      throw new UseAfterFreeError(this);
+    }
+    this.#rc++;
+    return this;
+  }
+  dispose() {
+    if (this.#rc <= 0) {
+      throw new UseAfterFreeError(this);
+    }
+    if (--this.#rc === 0) {
+      if (this.pval.isKnown) {
+        this.pval.val.dispose();
+      } else if (this.recipe) {
+        if (this.recipe.type === "Const") {
+          this.recipe.val.dispose();
+        } else if (this.recipe.type === "JaxprEqn") {
+          this.recipe.tracersIn.forEach((t) => t.dispose());
+        }
+      }
+    }
+  }
+  fullLower() {
+    if (this.pval.isKnown) {
+      const val = this.pval.val.ref;
+      this.dispose();
+      return val;
+    }
+    return this;
+  }
+};
+var PartialEvalTrace = class extends Trace {
+  newArg(pval) {
+    if (pval.isKnown) return new PartialEvalTracer(this, pval, null);
+    return new PartialEvalTracer(this, pval, { type: "LambdaBinding" });
+  }
+  pure(val) {
+    return new PartialEvalTracer(this, PartialVal.known(pureArray(val)), null);
+  }
+  lift = this.pure;
+  instantiateConst(tracer) {
+    if (!tracer.pval.isKnown) {
+      return tracer;
+    } else {
+      const pval = PartialVal.unknown(ShapedArray.fromAval(tracer.aval));
+      const val = tracer.pval.val.ref;
+      tracer.dispose();
+      return new PartialEvalTracer(this, pval, { type: "Const", val });
+    }
+  }
+  processPrimitive(primitive, tracers, params) {
+    if (tracers.every((t) => t.pval.isKnown)) {
+      return bind(
+        primitive,
+        tracers.map((t) => t.fullLower()),
+        params
+      );
+    }
+    if (primitive === "jit_call" /* JitCall */) {
+      return this.#partialEvalJaxpr(params.jaxpr, params.numConsts, tracers);
+    }
+    const tracersIn = tracers.map((t) => this.instantiateConst(t));
+    const avalsIn = tracersIn.map((t) => t.pval.aval);
+    const avalsOut = abstractEvalRules[primitive](avalsIn, params);
+    const recipe = {
+      type: "JaxprEqn",
+      prim: primitive,
+      tracersIn,
+      params,
+      avalsOut,
+      tracerRefsOut: []
+      // Populated later on
+    };
+    const tracersOut = avalsOut.map((aval, i) => {
+      if (i > 0) {
+        tracersIn.forEach((t) => t.ref);
+      }
+      return new PartialEvalTracer(this, PartialVal.unknown(aval), recipe);
+    });
+    recipe.tracerRefsOut = tracersOut.map((t) => new WeakRef(t));
+    return tracersOut;
+  }
+  /**
+   * Evaluate a Jaxpr on a set of PartialEvalTracers, computing as many known
+   * values as possible (with JIT) and forwarding the unknown ones.
+   *
+   * Used when encountering a JitCall rule during the trace.
+   */
+  #partialEvalJaxpr(jaxpr, numConsts, tracers) {
+    jaxpr = jaxpr.flatten();
+    const inUnknowns = tracers.map((t) => !t.pval.isKnown);
+    const { jaxpr1, jaxpr2, outUnknowns, numRes } = partialEvalJaxpr(
+      jaxpr,
+      inUnknowns
+    );
+    const [knownTracers, unknownTracers] = partitionList(inUnknowns, tracers);
+    const outs1Res = bind(
+      "jit_call" /* JitCall */,
+      knownTracers.map((t) => t.ref.fullLower()),
+      { jaxpr: jaxpr1, numConsts: 0 }
+    );
+    const outs1 = outs1Res.slice(0, jaxpr1.outs.length - numRes);
+    const res = outs1Res.slice(jaxpr1.outs.length - numRes);
+    const resTracers = res.map(
+      (x) => this.instantiateConst(fullRaise(this, x))
+    );
+    const recipe = {
+      type: "JaxprEqn",
+      prim: "jit_call" /* JitCall */,
+      tracersIn: resTracers.concat(unknownTracers),
+      params: { jaxpr: jaxpr2, numConsts: 0 },
+      avalsOut: jaxpr2.outs.map((x) => x.aval),
+      tracerRefsOut: []
+      // populated later
+    };
+    const outs2 = jaxpr2.outs.map(
+      (x) => new PartialEvalTracer(this, PartialVal.unknown(x.aval), recipe)
+    );
+    recipe.tracerRefsOut = outs2.map((t) => new WeakRef(t));
+    let i = 0;
+    let j = 0;
+    return outUnknowns.map((unk) => unk ? outs2[j++] : outs1[i++]);
+  }
+};
+function partialEvalJaxpr(jaxpr, inUnknowns, instantiate) {
+  jaxpr = jaxpr.flatten();
+  const knownIns = jaxpr.inBinders.filter((_, i) => !inUnknowns[i]);
+  const knownVars = new Set(knownIns);
+  const residuals = /* @__PURE__ */ new Set();
+  const eqns1 = [];
+  const eqns2 = [];
+  for (const eqn of jaxpr.eqns) {
+    if (eqn.primitive === "jit_call" /* JitCall */) {
+      throw new TypeError("partialEvalJaxpr requires flattened Jaxpr");
+    }
+    const hasUnknowns = eqn.inputs.some(
+      (x) => x instanceof Var && !knownVars.has(x)
+    );
+    if (hasUnknowns) {
+      for (const x of eqn.inputs) {
+        if (x instanceof Var && knownVars.has(x)) {
+          residuals.add(x);
+        }
+      }
+      eqns2.push(eqn);
+    } else {
+      eqns1.push(eqn);
+      for (const v of eqn.outBinders) {
+        knownVars.add(v);
+      }
+    }
+  }
+  const outUnknowns = jaxpr.outs.map(
+    (x) => x instanceof Var && !knownVars.has(x)
+  );
+  if (instantiate !== void 0) {
+    for (let i = 0; i < jaxpr.outs.length; i++) {
+      const x = jaxpr.outs[i];
+      if (instantiate[i] && !outUnknowns[i] && x instanceof Var) {
+        residuals.add(x);
+        outUnknowns[i] = true;
+      }
+    }
+  }
+  const residualsL = Array.from(residuals);
+  const [ins1, ins2] = partitionList(inUnknowns, jaxpr.inBinders);
+  const [outs1, outs2] = partitionList(outUnknowns, jaxpr.outs);
+  const jaxpr1 = new Jaxpr2(ins1, eqns1, outs1.concat(residualsL));
+  const jaxpr2 = new Jaxpr2(residualsL.concat(ins2), eqns2, outs2);
+  return { jaxpr1, jaxpr2, outUnknowns, numRes: residualsL.length };
+}
+function partialEvalGraphToJaxpr(tracersIn, tracersOut) {
+  const tracerToVar = /* @__PURE__ */ new Map();
+  const constToVar = /* @__PURE__ */ new Map();
+  const processedEqns = /* @__PURE__ */ new Set();
+  const eqns = [];
+  for (const t of tracersIn) {
+    tracerToVar.set(t, new Var(ShapedArray.fromAval(t.aval)));
+  }
+  for (const t of toposort(
+    tracersOut,
+    (t2) => t2.recipe?.type === "JaxprEqn" ? t2.recipe.tracersIn : []
+  )) {
+    if (!t.recipe) {
+      throw new TypeError("Tracer is missing a recipe, cannot construct Jaxpr");
+    }
+    if (t.recipe.type === "LambdaBinding") {
+      if (!tracersIn.includes(t)) {
+        throw new TypeError("LambdaBinding tracer not in input list");
+      }
+    } else if (t.recipe.type === "Const") {
+      const val = t.recipe.val;
+      let binder = constToVar.get(val);
+      if (!binder) {
+        binder = new Var(ShapedArray.fromAval(val.aval));
+        constToVar.set(val, binder);
+      }
+      tracerToVar.set(t, binder);
+    } else if (t.recipe.type === "JaxprEqn") {
+      if (!processedEqns.has(t.recipe)) {
+        processedEqns.add(t.recipe);
+        const tracersIn2 = t.recipe.tracersIn.map((t2) => tracerToVar.get(t2));
+        const outBinders = t.recipe.avalsOut.map((aval) => new Var(aval));
+        for (let i = 0; i < outBinders.length; i++) {
+          const tracerOut = t.recipe.tracerRefsOut[i].deref();
+          if (tracerOut) {
+            tracerToVar.set(tracerOut, outBinders[i]);
+          }
+        }
+        eqns.push(
+          new JaxprEqn(t.recipe.prim, tracersIn2, t.recipe.params, outBinders)
+        );
+      }
+    }
+  }
+  const [consts, constvars] = unzip2(constToVar.entries());
+  const inBinders = [
+    ...constvars,
+    ...tracersIn.map((t) => tracerToVar.get(t))
+  ];
+  const outVars = tracersOut.map((t) => tracerToVar.get(t));
+  const jaxpr = new Jaxpr2(inBinders, eqns, outVars);
+  typecheckJaxpr(jaxpr);
+  for (const t of consts) t.ref;
+  for (const t of tracersIn) t.dispose();
+  for (const t of tracersOut) t.dispose();
+  return { jaxpr: jaxpr.simplify(), consts };
+}
+var UndefPrimal = class {
+  aval;
+  constructor(aval) {
+    this.aval = ShapedArray.fromAval(aval);
+  }
+};
+function evalJaxprTransposed(jaxpr, args, cotangents) {
+  const knownPrimals = /* @__PURE__ */ new Map();
+  for (let i = 0; i < jaxpr.inBinders.length; i++) {
+    if (!(args[i] instanceof UndefPrimal)) {
+      knownPrimals.set(jaxpr.inBinders[i], args[i]);
+    }
+  }
+  const ctStore = /* @__PURE__ */ new Map();
+  const readCotangent = (v) => {
+    const ct = ctStore.get(v);
+    if (ct) {
+      ctStore.delete(v);
+      return ct;
+    } else {
+      return zeros(v.aval.shape, { dtype: v.aval.dtype });
+    }
+  };
+  const writeCotangent = (v, ct) => {
+    if (ct !== null) {
+      const oldCt = ctStore.get(v);
+      if (oldCt) ctStore.set(v, add(oldCt, ct));
+      else ctStore.set(v, ct);
+    }
+  };
+  for (let i = 0; i < jaxpr.outs.length; i++) {
+    const v = jaxpr.outs[i];
+    if (v instanceof Var) writeCotangent(v, cotangents[i]);
+  }
+  for (let i = jaxpr.eqns.length - 1; i >= 0; i--) {
+    const eqn = jaxpr.eqns[i];
+    const primalsIn = eqn.inputs.map(
+      (v) => v instanceof Lit ? scalar(v.value, { dtype: v.dtype }) : knownPrimals.has(v) ? knownPrimals.get(v).ref : new UndefPrimal(v.aval)
+    );
+    const cotangentsOut = eqn.outBinders.map(readCotangent);
+    const rule = transposeRules[eqn.primitive];
+    if (!rule) {
+      throw new TypeError(`Backward pass not implemented for ${eqn.primitive}`);
+    }
+    const cotangentsIn = rule(cotangentsOut, primalsIn, eqn.params);
+    for (let j = 0; j < eqn.inputs.length; j++) {
+      const v = eqn.inputs[j];
+      if (v instanceof Var && !knownPrimals.has(v)) {
+        writeCotangent(v, cotangentsIn[j]);
+      } else if (cotangentsIn[j] !== null) {
+        throw new Error("invariant violation: cotangent should be null");
+      }
+    }
+  }
+  for (const t of knownPrimals.values()) t.dispose();
+  const results = [];
+  for (let i = 0; i < jaxpr.inBinders.length; i++) {
+    if (args[i] instanceof UndefPrimal) {
+      results.push(readCotangent(jaxpr.inBinders[i]));
+    }
+  }
+  return results;
+}
+var NonlinearError = class extends TypeError {
+  constructor(primitive) {
+    super(`Nonlinear operation in backward pass for ${primitive}`);
+  }
+};
+var transposeRules = {
+  ["mul" /* Mul */]([ct], [x, y]) {
+    if (x instanceof UndefPrimal === y instanceof UndefPrimal)
+      throw new NonlinearError("mul" /* Mul */);
+    return [
+      x instanceof UndefPrimal ? mul(ct, y) : null,
+      y instanceof UndefPrimal ? mul(x, ct) : null
+    ];
+  },
+  ["neg" /* Neg */]([ct], [x]) {
+    if (!(x instanceof UndefPrimal)) throw new NonlinearError("neg" /* Neg */);
+    return [neg(ct)];
+  },
+  ["add" /* Add */]([ct], [x, y]) {
+    if (!(x instanceof UndefPrimal || y instanceof UndefPrimal))
+      throw new NonlinearError("add" /* Add */);
+    if (x instanceof UndefPrimal && y instanceof UndefPrimal)
+      return [ct.ref, ct];
+    return x instanceof UndefPrimal ? (y.dispose(), [ct, null]) : (x.dispose(), [null, ct]);
+  },
+  ["reduce_sum" /* ReduceSum */]([ct], [x], { axis }) {
+    if (!(x instanceof UndefPrimal))
+      throw new NonlinearError("reduce_sum" /* ReduceSum */);
+    return [broadcast(ct, x.aval.shape, axis)];
+  },
+  // BUG: Doesn't handle broadcasting.
+  ["where" /* Where */]([ct], [cond, x, y]) {
+    const cts = [null, null, null];
+    if (cond instanceof UndefPrimal) throw new NonlinearError("where" /* Where */);
+    if (x instanceof UndefPrimal) {
+      const zerosX = zeros(x.aval.shape, { dtype: x.aval.dtype });
+      cts[1] = where(cond.ref, ct.ref, zerosX);
+    } else {
+      x.dispose();
+    }
+    if (y instanceof UndefPrimal) {
+      const zerosY = zeros(x.aval.shape, { dtype: x.aval.dtype });
+      cts[2] = where(cond.ref, zerosY, ct.ref);
+    } else {
+      y.dispose();
+    }
+    ct.dispose();
+    cond.dispose();
+    return cts;
+  },
+  ["transpose" /* Transpose */]([ct], [x], { perm }) {
+    if (!(x instanceof UndefPrimal))
+      throw new NonlinearError("transpose" /* Transpose */);
+    return [transpose(ct, invertPermutation(perm))];
+  },
+  ["broadcast" /* Broadcast */]([ct], [x], { axis }) {
+    if (!(x instanceof UndefPrimal))
+      throw new NonlinearError("broadcast" /* Broadcast */);
+    return [reduceSum(ct, axis)];
+  },
+  ["reshape" /* Reshape */]([ct], [x], _) {
+    if (!(x instanceof UndefPrimal))
+      throw new NonlinearError("reshape" /* Reshape */);
+    return [reshape(ct, x.aval.shape)];
+  },
+  ["flip" /* Flip */]([ct], [x], { axis }) {
+    if (!(x instanceof UndefPrimal)) throw new NonlinearError("flip" /* Flip */);
+    return [flip(ct, axis)];
+  },
+  ["jit_call" /* JitCall */](cts, args, { jaxpr }) {
+    const undefPrimals = args.map((x) => x instanceof UndefPrimal);
+    const { newJaxpr, newConsts } = transposeJaxpr(jaxpr, undefPrimals);
+    const residuals = args.filter((x, i2) => !undefPrimals[i2]);
+    const outs = bind("jit_call" /* JitCall */, [...newConsts, ...residuals, ...cts], {
+      jaxpr: newJaxpr,
+      numConsts: newConsts.length
+    });
+    let i = 0;
+    return undefPrimals.map((isUndef) => isUndef ? outs[i++] : null);
+  }
+};
+var transposeJaxprCache = /* @__PURE__ */ new Map();
+function transposeJaxpr(jaxpr, undefPrimals) {
+  const cacheKey = JSON.stringify(undefPrimals);
+  const prevResult = transposeJaxprCache.get(jaxpr)?.get(cacheKey);
+  if (prevResult) return prevResult;
+  const { inTypes, outTypes } = typecheckJaxpr(jaxpr);
+  const forwardInTypes = inTypes.filter((_, i) => !undefPrimals[i]);
+  const { jaxpr: newJaxpr, consts: newConsts } = makeJaxpr(
+    (forwardIn, cotangents) => {
+      const args = [];
+      let forwardInIdx = 0;
+      for (let i = 0; i < undefPrimals.length; i++) {
+        if (undefPrimals[i]) args.push(new UndefPrimal(inTypes[i]));
+        else args.push(forwardIn[forwardInIdx++]);
+      }
+      return evalJaxprTransposed(jaxpr, args, cotangents);
+    }
+  )(forwardInTypes, outTypes);
+  typecheckJaxpr(newJaxpr);
+  const result = { newJaxpr, newConsts };
+  if (!transposeJaxprCache.has(jaxpr))
+    transposeJaxprCache.set(jaxpr, /* @__PURE__ */ new Map());
+  transposeJaxprCache.get(jaxpr).set(cacheKey, result);
+  return result;
+}
+function vjpFlat(f, primalsIn) {
+  const { primalsOut, jaxpr, consts } = linearizeFlatUtil(f, primalsIn);
+  const transposeInputs = [
+    ...consts.map((c) => c.ref),
+    // Explcitly list which arguments should be transposed.
+    ...primalsIn.map((t) => new UndefPrimal(t.aval))
+  ];
+  const fVjp = (...cotangents) => evalJaxprTransposed(jaxpr, transposeInputs, cotangents);
+  return [primalsOut, fVjp];
+}
+function vjp(f, ...primalsIn) {
+  const [primalsInFlat, inTree] = flatten(primalsIn);
+  const [fFlat, outTree] = flattenFun(f, inTree);
+  const [primalsOutFlat, fVjpFlat] = vjpFlat(
+    fFlat,
+    primalsInFlat.map(pureArray)
+  );
+  if (outTree.value === void 0) {
+    throw new Error("outTree was not set in vjp");
+  }
+  const primalsOut = unflatten(outTree.value, primalsOutFlat);
+  const fVjp = (cotangentsOut) => {
+    const [cotangentsOutFlat, outTree2] = flatten(cotangentsOut);
+    if (!outTree.value.equals(outTree2)) {
+      throw new TreeMismatchError("vjp", outTree.value, outTree2);
+    }
+    const cotangentsInFlat = fVjpFlat(...cotangentsOutFlat.map(pureArray));
+    return unflatten(inTree, cotangentsInFlat);
+  };
+  return [primalsOut, fVjp];
+}
+function grad(f) {
+  return (...x) => {
+    const [y, fVjp] = vjp(f, ...x);
+    if (!(y instanceof Tracer) || ndim(y) !== 0) {
+      throw new TypeError("grad requires a scalar output");
+    }
+    if (y.dtype !== "float32" /* Float32 */) {
+      throw new TypeError("grad currently only supports float32");
+    }
+    return fVjp(pureArray(1))[0];
+  };
+}
+function jacrev(f) {
+  return function jacobianReverse(x) {
+    if (x.shape.length !== 1) {
+      throw new TypeError("jacrev only supports 1D inputs");
+    }
+    const [size2] = x.shape;
+    const pullback = (ct) => vjp(f, x)[1](ct)[0];
+    return vmap(pullback, [1])(eye(size2, void 0, { dtype: x.dtype }));
+  };
+}
+
+// src/nn.ts
+var nn_exports = {};
+__export(nn_exports, {
+  identity: () => identity2,
+  logSigmoid: () => logSigmoid,
+  relu: () => relu,
+  relu6: () => relu6,
+  sigmoid: () => sigmoid,
+  silu: () => silu,
+  softSign: () => softSign,
+  softplus: () => softplus,
+  swish: () => swish
+});
+
+// src/numpy.ts
+var numpy_exports = {};
+__export(numpy_exports, {
+  Array: () => Array3,
+  DType: () => DType,
+  abs: () => abs,
+  absolute: () => absolute,
+  add: () => add2,
+  allclose: () => allclose,
+  arange: () => arange,
+  array: () => array,
+  bool: () => bool,
+  clip: () => clip,
+  complex64: () => complex64,
+  cos: () => cos2,
+  diag: () => diag,
+  diagonal: () => diagonal,
+  divide: () => divide,
+  dot: () => dot,
+  e: () => e,
+  equal: () => equal2,
+  eulerGamma: () => eulerGamma,
+  exp: () => exp2,
+  exp2: () => exp22,
+  eye: () => eye,
+  flip: () => flip2,
+  fliplr: () => fliplr,
+  flipud: () => flipud,
+  float32: () => float32,
+  full: () => full,
+  greater: () => greater2,
+  greaterEqual: () => greaterEqual2,
+  identity: () => identity,
+  inf: () => inf,
+  int32: () => int32,
+  less: () => less2,
+  lessEqual: () => lessEqual2,
+  linspace: () => linspace,
+  log: () => log2,
+  log10: () => log10,
+  log2: () => log22,
+  matmul: () => matmul,
+  matrixTranspose: () => matrixTranspose,
+  maximum: () => maximum,
+  meshgrid: () => meshgrid,
+  minimum: () => minimum,
+  moveaxis: () => moveaxis2,
+  multiply: () => multiply,
+  nan: () => nan,
+  ndim: () => ndim2,
+  negative: () => negative,
+  notEqual: () => notEqual2,
+  ones: () => ones,
+  permuteDims: () => permuteDims,
+  pi: () => pi,
+  ravel: () => ravel,
+  reciprocal: () => reciprocal2,
+  reshape: () => reshape2,
+  scalar: () => scalar,
+  shape: () => shape,
+  sin: () => sin2,
+  size: () => size,
+  square: () => square,
+  sum: () => sum,
+  tan: () => tan,
+  transpose: () => transpose2,
+  trueDivide: () => trueDivide,
+  trunc: () => trunc,
+  vdot: () => vdot,
+  vecdot: () => vecdot,
+  where: () => where2,
+  zeros: () => zeros
+});
+var float32 = "float32" /* Float32 */;
+var int32 = "int32" /* Int32 */;
+var bool = "bool" /* Bool */;
+var complex64 = "complex64" /* Complex64 */;
+var e = Math.E;
+var eulerGamma = 0.5772156649015329;
+var inf = Number.POSITIVE_INFINITY;
+var nan = Number.NaN;
+var pi = Math.PI;
+var add2 = add;
+var multiply = mul;
+var negative = neg;
+var reciprocal2 = reciprocal;
+var sin2 = sin;
+var cos2 = cos;
+var exp2 = exp;
+var log2 = log;
+var minimum = min;
+var maximum = max;
+var greater2 = greater;
+var less2 = less;
+var equal2 = equal;
+var notEqual2 = notEqual;
+var greaterEqual2 = greaterEqual;
+var lessEqual2 = lessEqual;
+var where2 = where;
+var transpose2 = transpose;
+var reshape2 = reshape;
+var sum = reduceSum;
+var moveaxis2 = moveaxis;
+var ndim2 = ndim;
+var shape = getShape;
+function size(a, axis) {
+  const s = shape(a);
+  return axis === void 0 ? prod(s) : s[axis];
+}
+function flip2(x, axis) {
+  const nd = ndim2(x);
+  if (axis === void 0) {
+    axis = range(nd);
+  } else if (typeof axis === "number") {
+    axis = [axis];
+  }
+  const seen = /* @__PURE__ */ new Set();
+  for (let i = 0; i < axis.length; i++) {
+    if (axis[i] >= nd || axis[i] < -nd) {
+      throw new TypeError(
+        `flip: axis ${axis[i]} out of bounds for array of ${nd} dimensions`
+      );
+    }
+    if (axis[i] < 0) axis[i] += nd;
+    if (seen.has(axis[i])) {
+      throw new TypeError(`flip: duplicate axis ${axis[i]} in axis list`);
+    }
+    seen.add(axis[i]);
+  }
+  return flip(x, axis);
+}
+function flipud(x) {
+  return flip2(x, 0);
+}
+function fliplr(x) {
+  return flip2(x, 1);
+}
+var permuteDims = transpose2;
+function ravel(a) {
+  return fudgeArray(a).ravel();
+}
+function diagonal(a, offset, axis1, axis2) {
+  return fudgeArray(a).diagonal(offset, axis1, axis2);
+}
+function matrixTranspose(x) {
+  const ar = fudgeArray(x);
+  if (ar.ndim < 2)
+    throw new TypeError("matrixTranspose only supports 2D+ arrays");
+  return ar.transpose([...range(ar.ndim - 2), ar.ndim - 1, ar.ndim - 2]);
+}
+function diag(v, k = 0) {
+  const a = fudgeArray(v);
+  if (!Number.isInteger(k))
+    throw new TypeError(`k must be an integer, got ${k}`);
+  if (a.ndim === 1) {
+    const n = a.shape[0];
+    const ret = where2(eye(n).equal(1), a, 0);
+    if (k !== 0) throw new Error("diag() for 1D arrays only for k=0");
+    return ret;
+  } else if (a.ndim === 2) {
+    return diagonal(a, k);
+  } else {
+    throw new TypeError("numpy.diag only supports 1D and 2D arrays");
+  }
+}
+function allclose(actual, expected, options) {
+  const { rtol = 1e-5, atol = 1e-8 } = options ?? {};
+  const x = array(actual);
+  const y = array(expected);
+  if (!deepEqual(x.shape, y.shape)) {
+    return false;
+  }
+  const xData = x.dataSync();
+  const yData = y.dataSync();
+  for (let i = 0; i < xData.length; i++) {
+    if (Math.abs(xData[i] - yData[i]) > atol + rtol * Math.abs(yData[i])) {
+      return false;
+    }
+  }
+  return true;
+}
+var matmul = jit(function matmul2(x, y) {
+  if (x.ndim === 0 || y.ndim === 0) {
+    throw new TypeError("matmul: x and y must be at least 1D");
+  }
+  if (y.ndim === 1) {
+    return x.mul(y).sum(x.ndim - 1);
+  }
+  x = x.reshape(x.shape.toSpliced(-1, 0, 1));
+  y = y.reshape(y.shape.toSpliced(-2, 0, 1)).transpose([
+    ...range(y.shape.length - 1),
+    y.shape.length,
+    y.shape.length - 1
+  ]);
+  return x.mul(y).sum(Math.max(x.ndim, y.ndim) - 1);
+});
+var dot = jit(function dot2(x, y) {
+  if (x.ndim === 0 || y.ndim === 0) {
+    return multiply(x, y);
+  }
+  if (y.ndim === 1) {
+    return x.mul(y).sum(x.ndim - 1);
+  }
+  x = x.reshape(x.shape.toSpliced(-1, 0, ...rep(y.ndim - 1, 1)));
+  y = y.transpose([
+    ...range(y.shape.length - 2),
+    y.shape.length - 1,
+    y.shape.length - 2
+  ]);
+  return x.mul(y).sum(x.ndim - 1);
+});
+var vecdot = jit(function vecdot2(x, y) {
+  return x.mul(y).sum(Math.max(x.ndim, y.ndim) - 1);
+});
+function vdot(x, y) {
+  return vecdot(ravel(x), ravel(y));
+}
+function meshgrid(xs, { indexing } = {}) {
+  indexing ??= "xy";
+  for (const x of xs) {
+    if (x.ndim !== 1) {
+      throw new TypeError(
+        `meshgrid: all inputs must be 1D arrays, got ${x.ndim}D array`
+      );
+    }
+  }
+  if (xs.length <= 1) return xs;
+  if (indexing === "xy") {
+    const [a, b, ...rest] = xs;
+    const [rb, ra, ...rrest] = meshgrid([b, a, ...rest], { indexing: "ij" });
+    return [ra, rb, ...rrest];
+  }
+  const shape2 = xs.map((x) => x.shape[0]);
+  return xs.map(
+    (x, i) => broadcast(x, shape2, [
+      ...range(i),
+      ...range(i + 1, xs.length)
+    ])
+  );
+}
+function clip(a, min2, max2) {
+  a = fudgeArray(a);
+  if (max2 !== void 0) {
+    a = minimum(a, max2);
+  }
+  if (min2 !== void 0) {
+    a = maximum(a, min2);
+  }
+  return a;
+}
+function absolute(x) {
+  x = fudgeArray(x);
+  return where2(less2(x.ref, 0), x.ref.mul(-1), x);
+}
+var abs = absolute;
+function square(x) {
+  x = fudgeArray(x);
+  return x.ref.mul(x);
+}
+function tan(x) {
+  x = fudgeArray(x);
+  return sin2(x.ref).div(cos2(x));
+}
+function trueDivide(x, y) {
+  x = fudgeArray(x);
+  y = fudgeArray(y);
+  if (!isFloatDtype(x.dtype) || !isFloatDtype(y.dtype)) {
+    throw new TypeError(
+      `trueDivide: x and y must be floating-point arrays, got ${x.dtype} and ${y.dtype}`
+    );
+  }
+  return x.div(y);
+}
+var divide = trueDivide;
+function trunc(x) {
+  return idiv(x, 1);
+}
+function exp22(p) {
+  return exp2(multiply(p, Math.LN2));
+}
+function log22(x) {
+  return log2(x).mul(Math.LOG2E);
+}
+function log10(x) {
+  return log2(x).mul(Math.LOG10E);
+}
+
+// src/nn.ts
+function relu(x) {
+  return maximum(x, 0);
+}
+function relu6(x) {
+  return clip(x, 0, 6);
+}
+function sigmoid(x) {
+  return reciprocal2(exp2(negative(x)).add(1));
+}
+function softplus(x) {
+  return log2(exp2(x).add(1));
+}
+function softSign(x) {
+  x = fudgeArray(x);
+  return x.ref.div(absolute(x).add(1));
+}
+function silu(x) {
+  x = fudgeArray(x);
+  return x.ref.mul(sigmoid(x));
+}
+var swish = silu;
+function logSigmoid(x) {
+  return negative(softplus(negative(x)));
+}
+var identity2 = fudgeArray;
+
+// src/polyfills.ts
+Symbol.dispose ??= Symbol.for("Symbol.dispose");
+Symbol.asyncDispose ??= Symbol.for("Symbol.asyncDispose");
+
+// src/index.ts
+var jvp2 = jvp;
+var vmap2 = vmap;
+var jacfwd2 = jacfwd;
+var makeJaxpr2 = makeJaxpr;
+var jit2 = jit;
+var linearize2 = linearize;
+var vjp2 = vjp;
+var grad2 = grad;
+var jacrev2 = jacrev;
+var jacobian = jacrev2;
+export {
+  devices,
+  grad2 as grad,
+  init,
+  jacfwd2 as jacfwd,
+  jacobian,
+  jacrev2 as jacrev,
+  jit2 as jit,
+  jvp2 as jvp,
+  linearize2 as linearize,
+  makeJaxpr2 as makeJaxpr,
+  nn_exports as nn,
+  numpy_exports as numpy,
+  setDevice,
+  tree_exports as tree,
+  vjp2 as vjp,
+  vmap2 as vmap
+};