Defining new Jax types with hijax
docs.jax.dev
Defining new Jax types with hijax
1–5 of 5 posts
Re: Defining new Jax types with hijax
#2> the usual tool is a pytree:
404.
Re: Defining new Jax types with hijax
#3Hijacking JAX with hijax. Such a neat name.
Re: Defining new Jax types with hijax
#4Why 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.
Re: Defining new Jax types with hijax
#5Why 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.
As I understand, hijax types can specify how Jax transformations apply to them, ie.e. how they are differentiated, vmapped, ...