Defining new Jax types with hijax

Defining new Jax types with hijax(docs.jax.dev)

14 points by fhchl 52 days ago | 4 comments

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.

fhchl 51 days ago | |

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

0xnyn 52 days ago |

Hijacking JAX with hijax. Such a neat name.

chrisjj 52 days ago |

> the usual tool is a pytree:

404.