jaxline

Experiment framework

Provides a Python-based framework for building distributed JAX training and evaluation experiments

GitHub

153 stars
13 watching
21 forks
Language: Python
last commit: about 1 year ago
Linked from 1 awesome list


Backlinks from these awesome lists:

Related projects:

Repository Description Stars
google-deepmind/tf2jax Converts TensorFlow functions to equivalent JAX Python functions. 109
google-deepmind/functa A repository containing code for a meta-learning experiment on image datasets 150
google-deepmind/chex A set of utilities for writing reliable JAX code 797
google-deepmind/jraph A lightweight library for working with graph neural networks in jax. 1,380
darshandeshpande/jax-models Provides a collection of deep learning models and utilities in JAX/Flax for research purposes. 151
google-research/visu3d An abstraction layer between various deep learning frameworks and your program. 149
sergiooramas/tartarus A Python module for Deep Learning experiments on Audio and Text data combining classification, recommendation, and matrix factorization techniques. 101
google/paxml A framework for configuring and running machine learning experiments on top of Jax. 461
google-deepmind/android_env A platform for defining reinforcement learning tasks on top of Android devices 1,029
google/jaxopt An open-source project providing hardware accelerated, batchable and differentiable optimizers in JAX for deep learning. 941
google-deepmind/einshape A unified reshaping library for JAX and other frameworks. 100
dwarfmaster/babel2 A framework for implementing agent-based experiments with cognitive processing and multi-agent interactions 6
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/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
luchris429/purejaxrl A high-performance implementation of reinforcement learning training pipelines using JAX and PyTorch-like functionality 755