Rose

icon by Aaron Weiss / CC BY 4.0

Rose is an automatic differentiation engine for the web, inspired by JAX.

Installation

With npm:

npm i rose

With Yarn:

yarn add rose

With pnpm:

pnpm add rose

With Bun:

bun add rose

Usage

This example defines custom derivatives for the builtin JavaScript logarithm and power functions, then computes the output, gradient, and Hessian for the power function applied with base 2 and exponent 3:

import { Dual, Real, Vec, add, compile, div, fn, mul, opaque, vjp } from "rose";

const log = opaque([Real], Real, Math.log);
log.jvp = fn([Dual], Dual, ({ re: x, du: dx }) => {
  return { re: log(x), du: div(dx, x) };
});

const pow = opaque([Real, Real], Real, Math.pow);
pow.jvp = fn([Dual, Dual], Dual, ({ re: x, du: dx }, { re: y, du: dy }) => {
  const z = pow(x, y);
  return { re: z, du: mul(add(mul(dx, div(y, x)), mul(dy, log(x))), z) };
});

const Vec2 = Vec(2, Real);
const Mat2 = Vec(2, Vec2);

const f = fn([Vec2], Real, ([x, y]) => pow(x, y));
const g = fn([Vec2], Vec2, (v) => vjp(f)(v).grad(1));
const h = fn([Vec2], Mat2, (v) => {
  const { grad } = vjp(g)(v);
  return [grad([1, 0]), grad([0, 1])];
});

const funcs = await Promise.all([compile(f), compile(g), compile(h)]);
console.log(funcs.map((func) => func([2, 3])));

With Vite

If you are using Vite then you will need to also install the vite-plugin-top-level-await package, because Rose internally uses top-level await, which Vite does not directly support. You must also include the following in your Vite config:

import { defineConfig } from "vite";
import topLevelAwait from "vite-plugin-top-level-await";

export default defineConfig({
  // the plugin described above
  plugins: [topLevelAwait()],

  // Vite bundles external dependencies by default in development mode, but that
  // process does not include assets; this option disables that particular kind
  // of bundling for Rose so that it can use its internal WebAssembly module
  optimizeDeps: { exclude: ["rose"] },
});

Contributing

See CONTRIBUTING.md.

License

Rose is licensed under the MIT License.