GitHub repository: shreyaspadhy/jaxutils
Primary language: Python
Description: Common utilities in JAX/Flax to use across research projects