paxml

ML framework

A framework for configuring and running machine learning experiments on top of Jax.

Pax is a Jax-based machine learning framework for training large scale models. Pax allows for advanced and fully configurable experimentation and parallelization, and has demonstrated industry leading model flop utilization rates.

GitHub

461 stars
15 watching
69 forks
Language: Python
last commit: almost 2 years ago
Linked from 1 awesome list

c4gptjaxlarge-language-modelsllmmodel-flopsparallelism

Backlinks from these awesome lists:

Related projects:

RepositoryDescriptionStars
google/praxisA library designed to be usable by other JAX-based ML projects, providing optimized layers for machine learning at scale.178
gomlx/gomlxAn accelerated machine learning framework for Go that provides tools and libraries for training, fine-tuning, and modifying machine learning models.553
erfanzar/easydelA framework for training and serving machine learning models using JAX, with a focus on flexibility and customization.212
symjax/symjaxA symbolic programming library for machine learning and deep learning applications.120
ai-hypercomputer/maxtextA high-performance LLM written in Python/Jax for training and inference on Google Cloud TPUs and GPUs.1,557
cdipaolo/gomlA Go library that provides an online machine learning framework with various algorithms and models1,574
google-deepmind/jaxlineProvides a Python-based framework for building distributed JAX training and evaluation experiments153
matthias-wright/flaxmodelsProvides pre-trained deep learning models for the Jax/Flax ecosystem.240
jaxgaussianprocesses/gpjaxProvides a low-level interface to Gaussian process models in JAX for flexible extension and customisation467
microsoft/0xdeca10bA framework for hosting and training machine learning models on a blockchain, enabling secure sharing and prediction without requiring users to pay for data or model updates.559
scicloj/scicloj.ml.clj-djlProvides pre-trained machine learning models for natural language processing tasks using Clojure and the clj-djl framework.0
google/objaxAn open-source machine learning framework designed to simplify and accelerate research772
rlouf/mcxTools and methods for Bayesian deep learning using probabilistic programming.328
neuraxio/neuraxleA machine learning pipeline library that enables the creation of modular and reusable data processing workflows610
yandex/repA toolset for building and running reproducible machine learning experiments in Python689