jax

There are 55 repositories under jax topic.

keras-team / keras
Deep Learning for humans
data-science deep-learning jax machine-learning neural-networks python pytorch tensorflow
Language:Python 63537
jax
jax-ml / jax
Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more
jax
Language:Python 33916
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 27279
lukasmasuch / best-of-ml-python
🏆 A ranked list of awesome machine learning Python libraries. Updated weekly.
automl chatgpt data-analysis data-science data-visualization data-visualizations deep-learning gpt gpt-3 jax keras machine-learning ml nlp python pytorch scikit-learn tensorflow transformer
22740
ivy-llc / ivy
Convert Machine Learning Code Between Frameworks
jax numpy python pytorch tensorflow
Language:Python 14235
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 10505
einops
arogozhnikov / einops
Flexible and powerful tensor operations for readable and reliable code (for pytorch, jax, TF and others)
chainer cupy deep-learning einops jax keras numpy pytorch tensor tensorflow
Language:Python 9263
google / trax
Trax — Deep Learning with Clear Code and Speed
deep-learning deep-reinforcement-learning jax machine-learning numpy reinforcement-learning transformer
Language:Python 8295
probml / pyprobml
Python code for "Probabilistic Machine learning" book by Kevin Murphy
blackjax colab flax jax jupyter-notebooks machine-learning numpyro pml probabilistic-programming pymc3 pyro pytorch tensorflow
Language:Jupyter Notebook 6938
google / flax
Flax is a neural network library for JAX that is designed for flexibility.
jax
Language:Jupyter Notebook 6898
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.
artificial-intelligence autograd bayesian-statistics convolutional-neural-networks data-science deep-learning ensemble-learning feature-extraction graduate-school information-theory interview-preparation jax jobs logistic-regression loss-functions machine-learning python pytorch pytorch-tutorial
4726
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 4638
tensorflow / datasets
TFDS is a collection of datasets ready to use with TensorFlow, Jax, ...
data dataset datasets jax machine-learning numpy tensorflow
Language:Python 4498
google-research / scenic
Scenic: A Jax Library for Computer Vision Research and Beyond
attention computer-vision deep-learning jax research transformers vision-transformer
Language:Python 3704
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 3594
iree-org / iree
A retargetable MLIR-based machine learning compiler and runtime toolkit.
compiler cuda jax machine-learning mlir onnx pytorch runtime spirv tensorflow vulkan
Language:C++ 3448
alpa-projects / alpa
Training and serving large-scale neural networks with auto parallelization.
alpa auto-parallelization compiler deep-learning distributed-computing distributed-training high-performance-computing jax llm machine-learning
Language:Python 3162
google-deepmind / dm-haiku
JAX-based neural network library
deep-learning deep-neural-networks jax machine-learning neural-networks
Language:Python 3113
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 3007
bethgelab / foolbox
A Python toolbox to create adversarial examples that fool neural networks in PyTorch, TensorFlow, and JAX
adversarial-attacks adversarial-examples jax keras machine-learning python pytorch tensorflow
Language:Python 2922
google / brax
Massively parallel rigidbody physics simulation on accelerator hardware.
jax physics-simulation reinforcement-learning robotics
Language:Jupyter Notebook 2917
NVIDIA / TransformerEngine
A library for accelerating Transformer models on NVIDIA GPUs, including using 8-bit and 4-bit floating point (FP8 and FP4) 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 fp4
Language:Python 2891
explosion / thinc
🔮 A refreshing functional take on deep learning, compatible with your favorite libraries
ai artificial-intelligence deep-learning functional-programming jax machine-learning machine-learning-library mxnet natural-language-processing nlp python pytorch spacy tensorflow type-checking
Language:Python 2878
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++ 2873
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 differentiable-computing quantum-chemistry qml jax python
Language:Python 2869
patrick-kidger / equinox
Elegant easy-to-use neural networks + scientific computing in JAX. https://docs.kidger.site/equinox/
deep-learning equinox jax neural-networks
Language:Python 2657
google-deepmind / mctx
Monte Carlo tree search in JAX
jax reinforcement-learning monte-carlo-tree-search
Language:Python 2554
pyro-ppl / numpyro
Probabilistic programming with NumPy powered by JAX for autograd and JIT compilation to GPU/TPU/CPU.
bayesian-inference hmc inference-algorithms jax mcmc numpy probabilistic-programming pyro
Language:Python 2546
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.
chatbot deep-learning flax jax language-model large-language-models llama natural-language-processing transformer
Language:Python 2499
google / neural-tangents
Fast and Easy Infinite Neural Networks in Python
bayesian-inference bayesian-networks deep-networks gaussian-processes gradient-descent gradient-flow infinite-networks jax kernel kernel-computation neural-networks neural-tangents training-dynamics
Language:Jupyter Notebook 2360
danijar / dreamerv3
Mastering Diverse Domains through World Models
artificial-intelligence general jax minecraft reinforcement-learning world-models
Language:Python 2342
apple / axlearn
An Extensible Deep Learning Library
deep-learning jax
Language:Python 2282
n2cholas / awesome-jax
JAX - A curated list of resources https://github.com/google/jax
autograd awesome awesome-list deep-learning jax machine-learning neural-network numpy xla
1962
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:Go 1956
AI-Hypercomputer / maxtext
A simple, performant and scalable Jax LLM!
deepseek fine-tuning gemma2 gemma3 gpt jax large-language-models llama2 llama3 llama4 llm mistral mixtral sft
Language:Python 1952
google-deepmind / penzai
A JAX research toolkit for building, editing, and visualizing neural networks.
fine-tuning interpretability jax neural-networks visualization
Language:Python 1824