Research
Parax: Parametric Modeling in JAX + Equinox [P]
**Paramax** (also referred to as 'Parax' in the Reddit post title) is a small Python library by Daniel Ward that provides parameterizations and parameter constraints for JAX PyTrees, designed to wo...
Paramax (also referred to as "Parax" in the Reddit post title) is a small Python library by Daniel Ward that provides parameterizations and parameter constraints for JAX PyTrees, designed to work alongside the Equinox neural network framework. It works by applying custom constraints or behaviors to PyTree components via "unwrappable placeholders," supporting use cases such as enforcing positivity, structured matrices, weight normalization, and marking parameters as non-trainable. The library is flexible enough to apply custom parameterizations to PyTrees from external libraries and is listed as a companion tool in the official Equinox ecosystem documentation.
Related
- [[p-ibu-boost-a-gbdt-library-where-splits-are-absolutely-rejec|[P] ibu-boost: a GBDT library where splits are absolutely rejected, not just relatively ranked[P]]]
- [[p-building-a-llm-from-scratch-with-mary-shelleys-frankenstei|[P] Building a LLM from scratch with Mary Shelley's 'Frankenstein' (on Kaggle)]]
- [[p-pca-before-truncation-makes-non-matryoshka-embeddings-comp|[P] PCA before truncation makes non-Matryoshka embeddings compressible: results on BGE-M3 [P]]]
- What if your HNSW index stored 3-bit embeddings instead of float32? [R]
Source: research