jax

There are 44 repositories under jax topic.

transformers
huggingface / transformers
🤗 Transformers: State-of-the-art Machine Learning for Pytorch, TensorFlow, and JAX.
nlp natural-language-processing pytorch language-model tensorflow bert language-models pytorch-transformers nlp-library transformer model-hub pretrained-models jax flax seq2seq speech-recognition hacktoberfest python machine-learning deep-learning
Language:Python 135833
keras-team / keras
Deep Learning for humans
data-science deep-learning jax machine-learning neural-networks python pytorch tensorflow
Language:Python 62149
jax
jax-ml / jax
Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more
jax
Language:Python 30644
huggingface / diffusers
🤗 Diffusers: State-of-the-art diffusion models for image and audio generation in PyTorch and FLAX.
deep-learning diffusion flax hacktoberfest image-generation image2image jax latent-diffusion-models pytorch score-based-generative-modeling stable-diffusion stable-diffusion-diffusers text2image
Language:Python 26449
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.
book computer-vision data-science deep-learning gaussian-processes hyperparameter-optimization jax kaggle keras machine-learning mxnet natural-language-processing notebook python pytorch recommender-system reinforcement-learning tensorflow
Language:Python 24126
ml-tooling / 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
18214
ivy-llc / ivy
Convert Machine Learning Code Between Frameworks
converter deep-learning ivy jax machine-learning neural-network numpy python pytorch tensorflow translation transpilation
Language:Python 14016
wandb
wandb / wandb
The AI developer platform. Use Weights & Biases to train and fine-tune models, and manage models from experimentation to production.
ai collaboration data-science data-versioning deep-learning experiment-track hyperparameter-optimization hyperparameter-search hyperparameter-tuning jax keras machine-learning ml-platform mlops model-versioning pytorch reinforcement-learning reproducibility tensorflow
Language:Python 9206
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 8557
google / trax
Trax — Deep Learning with Clear Code and Speed
jax numpy deep-learning deep-reinforcement-learning machine-learning transformer reinforcement-learning
Language:Python 8111
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 6566
google / flax
Flax is a neural network library for JAX that is designed for flexibility.
jax
Language:Jupyter Notebook 6170
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
4529
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 4456
tensorflow / datasets
TFDS is a collection of datasets ready to use with TensorFlow, Jax, ...
tensorflow machine-learning data datasets numpy jax dataset
Language:Python 4323
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 3345
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 3082
google-deepmind / dm-haiku
JAX-based neural network library
deep-learning deep-neural-networks jax machine-learning neural-networks
Language:Python 2911
iree-org / iree
A retargetable MLIR-based machine learning compiler and runtime toolkit.
mlir vulkan tensorflow spirv cuda jax pytorch compiler machine-learning runtime
Language:C++ 2870
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 2824
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 2787
lululxvi / deepxde
A library for scientific machine learning and physics-informed learning
deep-learning deeponet jax multi-fidelity-data neural-network operator paddle pde physics-informed-learning pinn pytorch scientific-machine-learning tensorflow
Language:Python 2787
grame-cncm / faust
Functional programming language for signal processing and sound synthesis
audio c c-plus-plus cmajor compiler cpp csharp dlang dsp faust functional-programming jax jsfx julia julia-language llvm rnbo rust wasm
Language:C++ 2595
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 flax jax optax pytorch pytorch-lightning tutorial tutorials uvadlc
Language:Jupyter Notebook 2592
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 2419
pennylane
PennyLaneAI / pennylane
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.
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 2384
google / brax
Massively parallel rigidbody physics simulation on accelerator hardware.
jax physics-simulation reinforcement-learning robotics
Language:Jupyter Notebook 2376
google-deepmind / mctx
Monte Carlo tree search in JAX
jax monte-carlo-tree-search reinforcement-learning
Language:Python 2366
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 2326
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 2285
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 2141
NVIDIA / TransformerEngine
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.
cuda deep-learning fp8 gpu jax machine-learning python pytorch
Language:Python 2004
apple / axlearn
An Extensible Deep Learning Library
deep-learning jax
Language:Python 1883
neonwatty / machine_learning_refined
Notes, examples, and Python demos for the 2nd edition of the university textbook "Machine Learning Refined" (published by Cambridge University Press).
artificial-intelligence autograd collab data-science deep-learning jax jupyter-notebook lecture-notes machine-learning machine-learning-algorithms mathematical-optimization neural-network numpy python slides
Language:Python 1713
google-deepmind / penzai
A JAX research toolkit for building, editing, and visualizing neural networks.
fine-tuning interpretability jax neural-networks visualization
Language:Python 1691
kubeflow / training-operator
Distributed ML 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 1625