chex
JAX helper library
A set of utilities for writing reliable JAX code
797 stars
17 watching
44 forks
Language: Python
last commit: 2 months ago
Linked from 2 awesome lists
n2cholas/awesome-jax
amrzv/awesome-colab-notebooksRelated projects:
| Repository | Description | Stars |
|---|---|---|
| google-deepmind/jraph | A lightweight library for working with graph neural networks in jax. | 1,380 |
| google-deepmind/einshape | A unified reshaping library for JAX and other frameworks. | 100 |
| google-deepmind/dm_pix | An image processing library built on top of JAX to provide optimized and parallelized functions for machine learning research. | 395 |
| google-deepmind/jaxline | Provides a Python-based framework for building distributed JAX training and evaluation experiments | 153 |
| google-deepmind/tf2jax | Converts TensorFlow functions to equivalent JAX Python functions. | 109 |
| google-deepmind/distrax | A library of probability distributions and bijectors with a focus on readability, extensibility, and compatibility with existing frameworks. | 538 |
 darshandeshpande/jax-models | Provides a collection of deep learning models and utilities in JAX/Flax for research purposes. | 151 |
 google/jaxopt | An open-source project providing hardware accelerated, batchable and differentiable optimizers in JAX for deep learning. | 941 |
| google-deepmind/optax | A gradient processing and optimization library designed to facilitate research and productivity in machine learning by providing building blocks for custom optimizers and gradient processing components. | 1,730 |
| google-deepmind/tree | A library for working with nested data structures in Python | 948 |
| google-deepmind/arnheim | A toolkit for generating collage images using neural networks and optimization algorithms | 235 |
 jgor/dork-cli | Tools for performing searches against a Google custom search engine and returning unique page results | 145 |
| google-deepmind/android_env | A platform for defining reinforcement learning tasks on top of Android devices | 1,029 |
 jaxgaussianprocesses/gpjax | Provides a low-level interface to Gaussian process models in JAX for flexible extension and customisation | 467 |
| google-deepmind/functa | A repository containing code for a meta-learning experiment on image datasets | 150 |