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: about 1 month ago
Linked from 1 awesome list

c4gptjaxlarge-language-modelsllmmodel-flopsparallelism

Backlinks from these awesome lists:

Related projects:

Repository Description Stars
google/praxis A library designed to be usable by other JAX-based ML projects, providing optimized layers for machine learning at scale. 178
gomlx/gomlx An accelerated machine learning framework for Go that provides tools and libraries for training, fine-tuning, and modifying machine learning models. 553
erfanzar/easydel A framework for training and serving machine learning models using JAX, with a focus on flexibility and customization. 212
symjax/symjax A symbolic programming library for machine learning and deep learning applications. 120
ai-hypercomputer/maxtext A high-performance LLM written in Python/Jax for training and inference on Google Cloud TPUs and GPUs. 1,557
cdipaolo/goml A Go library that provides an online machine learning framework with various algorithms and models 1,574
google-deepmind/jaxline Provides a Python-based framework for building distributed JAX training and evaluation experiments 153
matthias-wright/flaxmodels Provides pre-trained deep learning models for the Jax/Flax ecosystem. 240
jaxgaussianprocesses/gpjax Provides a low-level interface to Gaussian process models in JAX for flexible extension and customisation 467
microsoft/0xdeca10b A 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-djl Provides pre-trained machine learning models for natural language processing tasks using Clojure and the clj-djl framework. 0
google/objax An open-source machine learning framework designed to simplify and accelerate research 772
rlouf/mcx Tools and methods for Bayesian deep learning using probabilistic programming. 328
neuraxio/neuraxle A machine learning pipeline library that enables the creation of modular and reusable data processing workflows 610
yandex/rep A toolset for building and running reproducible machine learning experiments in Python 689