4 ms·
Defining new Jax types with hijax
- chrisjj 2mo ago> the usual tool is a pytree: 404.
- 0xnyn 2mo agoHijacking JAX with hijax. Such a neat name.
- peterdsharpe 2mo agoWhy 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.
- fhchl 2mo agoAs I understand, hijax types can specify how Jax transformations apply to them, ie.e. how they are differentiated, vmapped, ...