icon by Aaron Weiss / CC BY 4.0
Rose is an automatic differentiation engine for the web, inspired by JAX.
With npm:
npm i roseWith Yarn:
yarn add roseWith pnpm:
pnpm add roseWith Bun:
bun add roseThis 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])));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"] },
});See CONTRIBUTING.md.
Rose is licensed under the MIT License.