12 projects
Equimo
Implementation of modern multimodal models in JAX/Equinox
rollfast
JAX optimization primitives for experimental optimizers and schedulers.
justdata
A TensorFlow-native data pipeline library with a modality-neutral core and first-class computer vision and acoustic recipes
banax
Deep equilibrium models in JAX/Equinox
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)