jax-js progress

writing this from a borrowed desk in the corner of the office. cold morning, second coffee, someone's typing fast next to me.

the backlog is long — tabletop sims, the Jute thing, a couple of half-finished posts — but i want to keep momentum on jax-js for the morning. personal work feels like a stack sometimes: whatever i pushed last week is what i think about first. that's fine. exploration is supposed to be uneven.

progress

got jacfwd() working end-to-end this morning. the math part was easy; it was the type signatures that ate the time. the goal isn't to be perfectly precise — i want callers to feel like they're using a normal function, with arrays in and arrays out.

the types aren't fully sound. i don't think they need to be. what matters is that they work — autocomplete fires in the right places, mistakes light up red, and nothing actively lies.

// Walk a JsTree<A> and replace its A leaves with B, keeping shape intact.
type RemapLeaves<T, A, B> = T extends A
  ? B
  : T extends ReadonlyArray<infer U>
    ? RemapLeaves<U, A, B>[]
    : { [K in keyof T]: RemapLeaves<T[K], A, B> };

// Narrow a callable so its argument tuple is a subtype of T.
type ArgsExtend<F extends (...args: any[]) => any, T> =
  Parameters<F> extends T ? F : never;

/** Forward-mode JVP for an arbitrary JsTree input. */
export const jvp = core.jvp as <F extends (...args: any[]) => JsTree<NDArray>>(
  f: ArgsExtend<F, JsTree<NDArrayLike>>,
  primals: RemapLeaves<Parameters<F>, NDArray, NDArrayLike>,
  tangents: RemapLeaves<Parameters<F>, NDArray, NDArrayLike>,
) => [ReturnType<F>, ReturnType<F>];

/** Vectorize over a batched axis on one or more inputs. */
export const vmap = core.vmap as <F extends (...args: any[]) => JsTree<NDArray>>(
  f: ArgsExtend<F, JsTree<NDArrayLike>>,
  inAxes: RemapLeaves<Parameters<F>, NDArray, number>,
) => F;

/** Column-by-column Jacobian via forward-mode AD. */
export const jacfwd = core.jacfwd as <F extends (x: NDArray) => NDArray>(
  f: F,
  x: NDArray,
) => F;

vs code has been stalling every few minutes. could be the language server chewing through the conditional types. either way it's slowing me down more than the actual code is. going to restart everything and bolt on a few test helpers — at minimum an allclose() for arrays so i stop eyeballing diffs.

skipping the pretty-printer for now. it can wait.

— restarted. cleaner. on to tests.

today's thought

you can't argue someone into a different way of working. you can demonstrate, make change cheap, leave the door open — but the choice belongs to them. trying to shortcut that part usually costs more than the time it saves.

the underrated move is sitting down and actually listening. not as a tactic. as the thing.

back to the bugs

tests work now, which means the bugs are visible now. one of those rabbit holes where the symptom is twelve frames away from the cause. ended up learning that the place to hook in is at bind(), and that toString() has to exist for any of the diagnostics to be readable.

the root cause turned out to be a one-line typo inside fullLower(). an entire morning of confusion, deleted in a single character.

sachinkesiraju/jax-js·a3f1c0e

remove dispatch table from trace path

+84 −211·3h ago

noon

demos starting up. good crew.

on a team that ships quickly, you spend a lot of time half-understanding things. it's uncomfortable but probably correct — you're loading context faster than you can fully metabolize it. the trick is keeping a running list of what you don't get yet, instead of pretending you do.

putting an idea into words is most of the work. the rest is just typing.

protomaps is neat — serves whole-world maps via byte ranges out of a single file. paired with maplibre on the client you've got a credible mapbox gl replacement for free.