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: almost 3 years ago
Linked from 1 awesome list


Backlinks from these awesome lists:

Related projects:

RepositoryDescriptionStars
google-deepmind/tf2jaxConverts TensorFlow functions to equivalent JAX Python functions.109
google-deepmind/functaA repository containing code for a meta-learning experiment on image datasets150
google-deepmind/chexA set of utilities for writing reliable JAX code797
google-deepmind/jraphA lightweight library for working with graph neural networks in jax.1,380
darshandeshpande/jax-modelsProvides a collection of deep learning models and utilities in JAX/Flax for research purposes.151
google-research/visu3dAn abstraction layer between various deep learning frameworks and your program.149
sergiooramas/tartarusA Python module for Deep Learning experiments on Audio and Text data combining classification, recommendation, and matrix factorization techniques.101
google/paxmlA framework for configuring and running machine learning experiments on top of Jax.461
google-deepmind/android_envA platform for defining reinforcement learning tasks on top of Android devices1,029
google/jaxoptAn open-source project providing hardware accelerated, batchable and differentiable optimizers in JAX for deep learning.941
google-deepmind/einshapeA unified reshaping library for JAX and other frameworks.100
dwarfmaster/babel2A framework for implementing agent-based experiments with cognitive processing and multi-agent interactions6
google-deepmind/dm_pixAn image processing library built on top of JAX to provide optimized and parallelized functions for machine learning research.395
google-deepmind/optaxA 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/purejaxrlA high-performance implementation of reinforcement learning training pipelines using JAX and PyTorch-like functionality755