Awesome Open Source
Search
Programming Languages
Languages
All Categories
Categories
About
Search results for jax
jax
x
378 search results found
Transformers
⭐
127,491
🤗 Transformers: State-of-the-art Machine Learning for Pytorch, TensorFlow, and JAX.
Keras
⭐
61,137
Deep Learning for humans
Jax
⭐
27,067
Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more
D2l En
⭐
21,912
Interactive deep learning book with multi-framework code, math, and discussions. Adopted at 500 universities from 70 countries including Stanford, MIT, Harvard, and Cambridge.
Best Of Ml Python
⭐
14,990
🏆 A ranked list of awesome machine learning Python libraries. Updated weekly.
Ivy
⭐
13,904
The Unified AI Framework
Wandb
⭐
8,204
🔥 A tool for visualizing and tracking your machine learning experiments. This repo contains the CLI and Python API.
Trax
⭐
7,818
Trax — Deep Learning with Clear Code and Speed
Einops
⭐
7,548
Flexible and powerful tensor operations for readable and reliable code (for pytorch, jax, TF and others)
Pyprobml
⭐
6,066
Python code for "Probabilistic Machine learning" book by Kevin Murphy
Flax
⭐
5,144
Flax is a neural network library for JAX that is designed for flexibility.
Datasets
⭐
4,094
TFDS is a collection of datasets ready to use with TensorFlow, Jax, ...
Whisper Jax
⭐
3,824
JAX implementation of OpenAI's Whisper model for up to 70x speed-up on TPU.
Interviews.ai
⭐
3,146
It is my belief that you, the postgraduate students and job-seekers for whom the book is primarily meant will benefit from reading it; however, it is my hope that even the most experienced researchers will find it fascinating as well.
Alpa
⭐
2,878
Training and serving large-scale neural networks with auto parallelization.
Dm Haiku
⭐
2,807
JAX-based neural network library
Thinc
⭐
2,774
🔮 A refreshing functional take on deep learning, compatible with your favorite libraries
Scenic
⭐
2,733
Scenic: A Jax Library for Computer Vision Research and Beyond
Foolbox
⭐
2,600
A Python toolbox to create adversarial examples that fool neural networks in PyTorch, TensorFlow, and JAX
Iree
⭐
2,418
A retargetable MLIR-based machine learning compiler and runtime toolkit.
Faust
⭐
2,351
Functional programming language for signal processing and sound synthesis
Mctx
⭐
2,201
Monte Carlo tree search in JAX
Pennylane
⭐
2,165
PennyLane is a cross-platform Python library for quantum computing, quantum machine learning, and quantum chemistry. Train a quantum computer the same way as a neural network.
Neural Tangents
⭐
2,149
Fast and Easy Infinite Neural Networks in Python
Deepxde
⭐
2,135
A library for scientific machine learning and physics-informed learning
Easylm
⭐
2,087
Large language models (LLMs) made easy, EasyLM is a one stop solution for pre-training, finetuning, evaluating and serving LLMs in JAX/Flax.
Numpyro
⭐
1,966
Probabilistic programming with NumPy powered by JAX for autograd and JIT compilation to GPU/TPU/CPU.
Brax
⭐
1,929
Massively parallel rigidbody physics simulation on accelerator hardware.
Uvadlc_notebooks
⭐
1,860
Repository of Jupyter notebook tutorials for teaching the Deep Learning Course at the University of Amsterdam (MSc AI), Fall 2022/Spring 2022
Equinox
⭐
1,614
Elegant easy-to-use neural networks + scientific computing in JAX. https://docs.kidger.site/equinox/
Machine_learning_refined
⭐
1,574
Notes, examples, and Python demos for the 2nd edition of the textbook "Machine Learning Refined" (published by Cambridge University Press).
Tensorly
⭐
1,466
TensorLy: Tensor Learning in Python.
Training Operator
⭐
1,447
Distributed ML Training and Fine-Tuning on Kubernetes
Big_transfer
⭐
1,421
Official repository for the "Big Transfer (BiT): General Visual Representation Learning" paper.
Jraph
⭐
1,233
A Graph Neural Network Library in Jax
Transformerengine
⭐
1,172
A library for accelerating Transformer models on NVIDIA GPUs, including using 8-bit floating point (FP8) precision on Hopper and Ada GPUs, to provide better performance with lower memory utilization in both training and inference.
Awesome Jax
⭐
1,156
JAX - A curated list of resources https://github.com/google/jax
Jraph
⭐
1,141
A Graph Neural Network Library in Jax
Diffrax
⭐
1,131
Numerical differential equation solvers in JAX. Autodifferentiable and GPU-capable. https://docs.kidger.site/diffrax/
Jaxopt
⭐
857
Hardware accelerated, batchable and differentiable optimizers in JAX.
Fortuna
⭐
824
A Library for Uncertainty Quantification.
Score_sde
⭐
782
Official code for Score-Based Generative Modeling through Stochastic Differential Equations (ICLR 2021, Oral)
Jaxtyping
⭐
759
Type annotations and runtime checking for shape and dtype of JAX/NumPy/PyTorch/etc. arrays. https://docs.kidger.site/jaxtyping/
Dawdreamer
⭐
750
Digital Audio Workstation with Python; VST instruments/effects, parameter automation, FAUST, JAX, Warp Markers, and JUCE processors
Dreamerv3
⭐
749
Mastering Diverse Domains through World Models
Mosec
⭐
661
A high-performance ML model serving framework, offers dynamic batching and CPU/GPU pipelines to fully exploit your compute machine
Eagerpy
⭐
645
PyTorch, TensorFlow, JAX and NumPy — all of them natively using the same code
Long Range Arena
⭐
635
Long Range Arena for Benchmarking Efficient Transformers
Tensor Sensor
⭐
634
The goal of this library is to generate more helpful exception messages for matrix algebra expressions for numpy, pytorch, jax, tensorflow, keras, fastai.
Mava
⭐
593
🦁 A research-friendly codebase for fast experimentation of multi-agent reinforcement learning in JAX
Jaxrl
⭐
585
JAX (Flax) implementation of algorithms for Deep Reinforcement Learning with continuous action spaces.
Prompt Tuning
⭐
580
Original Implementation of Prompt Tuning from Lester, et al, 2021
Treasure Of Transformers
⭐
541
💁 Awesome Treasure of Transformers Models for Natural Language processing contains papers, videos, blogs, official repo along with colab Notebooks. 🛫☑️
Dynamax
⭐
540
State Space Models library in JAX
Jax Cfd
⭐
527
Computational Fluid Dynamics in JAX
Pdebench
⭐
522
PDEBench: An Extensive Benchmark for Scientific Machine Learning
Ml Aim
⭐
515
This repository provides the code and model checkpoints of the research paper: Scalable Pre-training of Large Autoregressive Image Models
Get Started With Jax
⭐
487
The purpose of this repo is to make it easy to get started with JAX, Flax, and Haiku. It contains my "Machine Learning with JAX" series of tutorials (YouTube videos and Jupyter Notebooks) as well as the content I found useful while learning about the JAX ecosystem.
Fast Soft Sort
⭐
476
Fast Differentiable Sorting and Ranking
Netket
⭐
476
Machine learning algorithms for many-body quantum systems
Purejaxrl
⭐
460
Really Fast End-to-End Jax RL Implementations
Jumanji
⭐
439
🕹️ A diverse suite of scalable reinforcement learning environments in JAX
Neuralcompression
⭐
430
A collection of tools for neural compression enthusiasts.
Elegy
⭐
430
A High Level API for Deep Learning in JAX
Ott
⭐
417
Optimal Transport tools implemented with the JAX framework, to get auto-diff, parallel and jit-able computations.
Dataloader
⭐
390
The merlin dataloader lets you rapidly load tabular data for training deep leaning models with TensorFlow, PyTorch or JAX
Tpu Starter
⭐
381
Everything you want to know about Google Cloud TPU
Git Re Basin
⭐
378
Code release for "Git Re-Basin: Merging Models modulo Permutation Symmetries"
Machine Learning Guide
⭐
366
Machine learning Guide. Learn all about Machine Learning Tools, Libraries, Frameworks, Large Language Models (LLMs), and Training Models.
Gpjax
⭐
352
Gaussian processes in JAX.
Dm_pix
⭐
352
PIX is an image processing library in JAX, for JAX.
Mpi4jax
⭐
348
Zero-copy MPI communication of JAX arrays, for turbo-charged HPC applications in Python ⚡
Paxml
⭐
335
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.
Tevatron
⭐
334
Tevatron - A flexible toolkit for neural retrieval research and development.
Autobound
⭐
325
AutoBound automatically computes upper and lower bounds on functions.
Veros
⭐
307
The versatile ocean simulator, in pure Python, powered by JAX.
Learning Python Physics Informed Machine Learning Pinns Deeponets
⭐
297
Physics Informed Machine Learning Tutorials (Pytorch and Jax)
Extending Jax
⭐
293
Extending JAX with custom C++ and CUDA code
Pyhpc Benchmarks
⭐
290
A suite of benchmarks for CPU and GPU performance of the most popular high-performance libraries for Python 🚀
Scimlbenchmarks.jl
⭐
279
Scientific machine learning (SciML) benchmarks, AI for science, and (differential) equation solvers. Covers Julia, Python (PyTorch, Jax), MATLAB, R
Opennsfw2
⭐
272
Keras implementation of the Yahoo Open-NSFW model
Sympy2jax
⭐
272
Turn SymPy expressions into trainable JAX expressions.
Squirrel Core
⭐
271
A Python library that enables ML teams to share, load, and transform data in a collaborative, flexible, and efficient way 🌰
Pyhf
⭐
265
pure-Python HistFactory implementation with tensors and autodiff
Skrl
⭐
263
Modular reinforcement learning library (on PyTorch and JAX) with support for NVIDIA Isaac Gym, Isaac Orbit and Omniverse Isaac Gym
Commonlooputils
⭐
261
CLU lets you write beautiful training loops in JAX.
Lineax
⭐
261
Linear solvers in JAX and Equinox. https://docs.kidger.site/lineax
Pgx
⭐
260
🎲 Vectorized RL game environments in JAX [NeurIPS23]
Gradcache
⭐
254
Run Effective Large Batch Contrastive Learning Beyond GPU/TPU Memory Constraint
Fedjax
⭐
240
FedJAX is a JAX-based open source library for Federated Learning simulations that emphasizes ease-of-use in research.
Qdax
⭐
233
Accelerated Quality-Diversity
Annotated S4
⭐
231
Implementation of https://srush.github.io/annotated-s4
Tensorcircuit
⭐
230
Tensor network based quantum software framework for the NISQ era
Funsor
⭐
224
Functional tensors for probabilistic programming
Treex
⭐
217
A Pytree Module system for Deep Learning in JAX
Jaxfluids
⭐
203
Differentiable Fluid Dynamics Package
L2p
⭐
202
Learning to Prompt (L2P) for Continual Learning @ CVPR22 and DualPrompt: Complementary Prompting for Rehearsal-free Continual Learning @ ECCV22
Flaxmodels
⭐
201
Pretrained deep learning models for Jax/Flax: StyleGAN2, GPT2, VGG, ResNet, etc.
Bayesnewton
⭐
196
Bayes-Newton—A Gaussian process library in JAX, with a unifying view of approximate Bayesian inference as variants of Newton's method.
Jax Am
⭐
194
Additive manufacturing simulation with JAX.
1-100 of 378 search results
Next >
Privacy
|
About
|
Terms
|
Follow Us On Twitter
Copyright 2018-2024 Awesome Open Source. All rights reserved.