jax-js progress

feb 8 — at val

currently doing this at the val town office, hosted by justin-b. it’s a chilly day

on the backlog: tabletop.js, and Jute. lots of stuff, but let me make progress on this for the morning! sometimes feels like my personal work is a stack: recency bias toward what I work on. but that’s fine, exploration is hard work. let’s have fun

progress

I just implemented jacfwd(). got distracted by type signatures a bit, but I think things are going roughly smoothly. some folks next to me are collaborating on an API.

one important thing to remember is that there’s some type signatures that aren’t quite exactly right, but what matters is that they work

// Convert a subtype of JsTree<A> into a JsTree<B>, with the same structure.
type MapJsTree<T, A, B> = T extends A
  ? B
  : T extends globalThis.Array<infer U>
    ? MapJsTree<U, A, B>[]
    : { [K in keyof T]: MapJsTree<T[K], A, B> };

// Assert that a function's arguments are a subtype of the given type.
type WithArgsSubtype<F extends (args: any[]) => any, T> =
  Parameters<F> extends T ? F : never;

/** Compute the forward-mode Jacobian-vector product for a function. */
export const jvp = core.jvp as <F extends (...args: any[]) => JsTree<Array>>(
  f: WithArgsSubtype<F, JsTree<ArrayLike>>,
  primals: MapJsTree<Parameters<F>, Array, ArrayLike>,
  tangents: MapJsTree<Parameters<F>, Array, ArrayLike>
) => [ReturnType<F>, ReturnType<F>];

/** Vectorize an operation on a batched axis for one or more inputs. */
export const vmap = core.vmap as <F extends (...args: any[]) => JsTree<Array>>(
  f: WithArgsSubtype<F, JsTree<ArrayLike>>,
  inAxes: MapJsTree<Parameters<F>, Array, number>
) => F;

/** Compute the Jacobian evaluated column-by-column by forward-mode AD. */
export const jacfwd = core.jacfwd as <F extends (x: Array) => Array>(
  f: F,
  x: Array
) => F;

i’m feeling a bit annoyed by vs code freezing every so often. as well as my lack of tests. let me restart my computer for perf, and then I can add some utilities to assert arrays are close to each other, which will be needed for tests.

I’ll skip on the utility for printing arrays for now — we’ll get there later!

—my computer has restarted.

today’s idea

you cannot tell other people how to work. you can only show them, make it easy to change, give suggestions. but it’s ultimately their choice, and learning how to work together with others takes time. relationships are damaged when we aren’t patient and generous with each other.

sometimes, you just need to sit down with the other party, and hear them out. there should always be a place for honest listening

…and back to tests

i got tests to work, and then my tests promptly discovered several bugs (lol) and then in the process of debugging futily, I learned the core place where you need to intercept stuff at bind(), and then the necessity of having a toString() LOL.

but yeah it turns out a one-line typo in fullLower() caused all this mess. now you know — experiences!

and now it’s noon

so they’re doing demos. a nice group of people

when you’re working on a team that works fast, you often find yourself in a situation where you don’t understand everything. a primitive capability. context, understanding, some representation. design docs take energy to write

express your ideas is a lot of the work, perhaps it’s the productive part

protomaps! a nice file format for serving full-world maps via byte ranges. seems like a cool system. maplibre is the client format. replace mapbox gl :)