jax

There are 52 repositories under jax topic.

keras-team / keras
Deep Learning for humans
deep-learning tensorflow neural-networks machine-learning data-science python jax pytorch
Language:Python 63406
jax
jax-ml / jax
Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more
jax
Language:Python 33440
d2l-ai / d2l-en
Interactive deep learning book with multi-framework code, math, and discussions. Adopted at 500 universities from 70 countries including Stanford, MIT, Harvard, and Cambridge.
deep-learning machine-learning book notebook computer-vision natural-language-processing python kaggle data-science mxnet pytorch tensorflow keras gaussian-processes hyperparameter-optimization recommender-system reinforcement-learning jax
Language:Python 26843
lukasmasuch / best-of-ml-python
🏆 A ranked list of awesome machine learning Python libraries. Updated weekly.
python machine-learning data-science nlp data-visualization tensorflow pytorch scikit-learn keras data-visualizations deep-learning automl ml jax transformer data-analysis chatgpt gpt gpt-3
21916
ivy-llc / ivy
Convert Machine Learning Code Between Frameworks
python tensorflow pytorch numpy jax
Language:Python 14246
wandb
wandb / wandb
The AI developer platform. Use Weights & Biases to train and fine-tune models, and manage models from experimentation to production.
machine-learning experiment-track deep-learning keras tensorflow pytorch hyperparameter-search reinforcement-learning mlops data-science collaboration hyperparameter-optimization reproducibility hyperparameter-tuning data-versioning model-versioning ml-platform jax ai
Language:Python 10320
einops
arogozhnikov / einops
Flexible and powerful tensor operations for readable and reliable code (for pytorch, jax, TF and others)
deep-learning pytorch tensorflow numpy cupy chainer keras tensor jax einops
Language:Python 9167
google / trax
Trax — Deep Learning with Clear Code and Speed
jax numpy deep-learning deep-reinforcement-learning machine-learning transformer reinforcement-learning
Language:Python 8279
probml / pyprobml
Python code for "Probabilistic Machine learning" book by Kevin Murphy
jupyter-notebooks machine-learning colab pml probabilistic-programming jax tensorflow pytorch pymc3 numpyro flax pyro blackjax
Language:Jupyter Notebook 6887
google / flax
Flax is a neural network library for JAX that is designed for flexibility.
jax
Language:Jupyter Notebook 6802
BoltzmannEntropy / interviews.ai
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.
data-science machine-learning deep-learning interview-preparation jobs artificial-intelligence pytorch-tutorial graduate-school pytorch python jax autograd information-theory bayesian-statistics convolutional-neural-networks ensemble-learning feature-extraction logistic-regression loss-functions
4693
sanchit-gandhi / whisper-jax
JAX implementation of OpenAI's Whisper model for up to 70x speed-up on TPU.
deep-learning jax speech-recognition speech-to-text whisper
Language:Jupyter Notebook 4633
tensorflow / datasets
TFDS is a collection of datasets ready to use with TensorFlow, Jax, ...
tensorflow machine-learning data datasets numpy jax dataset
Language:Python 4481
google-research / scenic
Scenic: A Jax Library for Computer Vision Research and Beyond
jax computer-vision deep-learning research attention transformers vision-transformer
Language:Python 3661
lululxvi / deepxde
A library for scientific machine learning and physics-informed learning
neural-network deep-learning scientific-machine-learning pinn multi-fidelity-data operator pytorch physics-informed-learning jax deeponet paddle pde tensorflow
Language:Python 3477
iree-org / iree
A retargetable MLIR-based machine learning compiler and runtime toolkit.
mlir vulkan tensorflow spirv cuda jax pytorch compiler machine-learning runtime onnx
Language:C++ 3347
alpa-projects / alpa
Training and serving large-scale neural networks with auto parallelization.
deep-learning machine-learning compiler distributed-training high-performance-computing alpa jax distributed-computing llm auto-parallelization
Language:Python 3153
google-deepmind / dm-haiku
JAX-based neural network library
machine-learning neural-networks jax deep-learning deep-neural-networks
Language:Python 3090
phlippe / uvadlc_notebooks
Repository of Jupyter notebook tutorials for teaching the Deep Learning Course at the University of Amsterdam (MSc AI), Fall 2023
deep-learning tutorials uvadlc pytorch pytorch-lightning tutorial flax jax optax
Language:Jupyter Notebook 2973
bethgelab / foolbox
A Python toolbox to create adversarial examples that fool neural networks in PyTorch, TensorFlow, and JAX
adversarial-examples machine-learning python adversarial-attacks pytorch tensorflow jax keras
Language:Python 2898
explosion / thinc
🔮 A refreshing functional take on deep learning, compatible with your favorite libraries
machine-learning deep-learning artificial-intelligence ai python nlp natural-language-processing spacy machine-learning-library type-checking functional-programming pytorch tensorflow mxnet jax
Language:Python 2877
google / brax
Massively parallel rigidbody physics simulation on accelerator hardware.
jax reinforcement-learning robotics physics-simulation
Language:Jupyter Notebook 2854
grame-cncm / faust
Functional programming language for signal processing and sound synthesis
faust dsp audio compiler wasm c cpp llvm functional-programming rust dlang c-plus-plus csharp julia julia-language jax cmajor jsfx rnbo
Language:C++ 2831
pennylane
PennyLaneAI / pennylane
PennyLane is a cross-platform Python library for quantum computing, quantum machine learning, and quantum chemistry. Built by researchers, for research.
quantum machine-learning deep-learning neural-network optimization quantum-computing quantum-machine-learning automatic-differentiation tensorflow pytorch autograd qiskit cirq strawberryfields differentiable-computing quantum-chemistry qml jax python hacktoberfest
Language:Python 2811
NVIDIA / TransformerEngine
A library for accelerating Transformer models on NVIDIA GPUs, including using 8-bit floating point (FP8) precision on Hopper, Ada and Blackwell GPUs, to provide better performance with lower memory utilization in both training and inference.
cuda deep-learning gpu machine-learning python pytorch fp8 jax
Language:Python 2727
google-deepmind / mctx
Monte Carlo tree search in JAX
jax reinforcement-learning monte-carlo-tree-search
Language:Python 2532
pyro-ppl / numpyro
Probabilistic programming with NumPy powered by JAX for autograd and JIT compilation to GPU/TPU/CPU.
pyro jax hmc inference-algorithms numpy bayesian-inference probabilistic-programming mcmc
Language:Python 2527
patrick-kidger / equinox
Elegant easy-to-use neural networks + scientific computing in JAX. https://docs.kidger.site/equinox/
deep-learning jax neural-networks equinox
Language:Python 2515
young-geng / EasyLM
Large language models (LLMs) made easy, EasyLM is a one stop solution for pre-training, finetuning, evaluating and serving LLMs in JAX/Flax.
deep-learning flax jax language-model natural-language-processing transformer large-language-models chatbot llama
Language:Python 2499
google / neural-tangents
Fast and Easy Infinite Neural Networks in Python
neural-networks infinite-networks gaussian-processes training-dynamics neural-tangents jax kernel deep-networks kernel-computation gradient-descent bayesian-networks bayesian-inference gradient-flow
Language:Jupyter Notebook 2357
apple / axlearn
An Extensible Deep Learning Library
deep-learning jax
Language:Python 2246
danijar / dreamerv3
Mastering Diverse Domains through World Models
artificial-intelligence general minecraft reinforcement-learning world-models jax
Language:Python 2131
n2cholas / awesome-jax
JAX - A curated list of resources https://github.com/google/jax
awesome awesome-list jax machine-learning neural-network deep-learning xla autograd numpy
1918
kubeflow / trainer
Distributed AI Model Training and Fine-Tuning on Kubernetes
ai distributed fine-tuning gpu huggingface jax kubeflow kubernetes llm machine-learning mlops python pytorch tensorflow xgboost
Language:Python 1909
AI-Hypercomputer / maxtext
A simple, performant and scalable Jax LLM!
large-language-models llm gpt deepseek fine-tuning gemma2 gemma3 jax llama2 llama3 llama4 mistral mixtral sft
Language:Python 1900
google-deepmind / penzai
A JAX research toolkit for building, editing, and visualizing neural networks.
fine-tuning interpretability jax neural-networks visualization
Language:Python 1818