11 projects
banax
Deep equilibrium models in JAX/Equinox
Equimo
Implementation of popular vision models in Jax
rollfast
JAX implementation of experimental optimizers and schedulers.
psgdax
JAX implementation of the Preconditioned Stochastic Gradient Descent (PSGD) optimizer.
anchorax
Fixed-point solvers and differentiation utilities for implicit neural networks in JAX/Equinox.
hsf
A simple yet exhaustive segmentation tool of the Hippocampal Subfields in T1w and T2w MRIs.
roiloc
A simple package to center and crop T1w & T2w MRIs around a given region of interest by its name.
jamp
JAMP is a Mixed Precision library for JAX. Forked from DeepMind's JMP.
jimmy-vision
Implementations of popular vision models in Jax (Flax NNX)
