GitHub repository: soraros/equinox
Primary language: Python
Description: Callable PyTrees and filtered JIT/grad transformations => neural networks in JAX.