Defining new Jax types with hijax

14 points

1/21/1970

a day ago

by fhchl

Comments

peterdsharpe

Why use this over Equinox's Module, plus jaxtyping? That seems like the cleaner (and more established) solution, but maybe there are pros/cons I'm not seeing.

a day ago

fhchl

As I understand, hijax types can specify how Jax transformations apply to them, ie.e. how they are differentiated, vmapped, ...

8 hours ago

0xnyn

Hijacking JAX with hijax. Such a neat name.

a day ago

chrisjj

> the usual tool is a pytree:

404.

a day ago