jax

There are 38 repositories under jax topic.

transformers
huggingface / transformers
🤗 Transformers: State-of-the-art Machine Learning for Pytorch, TensorFlow, and JAX.
bert deep-learning flax hacktoberfest jax language-model language-models machine-learning model-hub natural-language-processing nlp nlp-library pretrained-models python pytorch pytorch-transformers seq2seq speech-recognition tensorflow transformer
Language:Python 126094
keras-team / keras
Deep Learning for humans
deep-learning tensorflow neural-networks machine-learning data-science python jax pytorch
Language:Python 61031
jax
google / jax
Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more
jax
Language:Python 28174
huggingface / diffusers
🤗 Diffusers: State-of-the-art diffusion models for image and audio generation in PyTorch and FLAX.
deep-learning diffusion image-generation pytorch score-based-generative-modeling image2image text2image stable-diffusion stable-diffusion-diffusers hacktoberfest flax jax latent-diffusion-models
Language:Python 22846
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 21905
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
15665
unifyai / ivy
The Unified AI Framework
python machine-learning deep-learning neural-network gpu autograd ivy abstraction template tensorflow pytorch mxnet numpy jax
Language:Python 14023
wandb
wandb / wandb
🔥 A tool for visualizing and tracking your machine learning experiments. This repo contains the CLI and Python API.
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
Language:Python 8308
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 7964
google / trax
Trax — Deep Learning with Clear Code and Speed
jax numpy deep-learning deep-reinforcement-learning machine-learning transformer reinforcement-learning
Language:Python 7962
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 6278
google / flax
Flax is a neural network library for JAX that is designed for flexibility.
jax
Language:Python 5561
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
4441
tensorflow / datasets
TFDS is a collection of datasets ready to use with TensorFlow, Jax, ...
data dataset datasets jax machine-learning numpy tensorflow
Language:Python 4198
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 4119
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 3035
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 2989
google-deepmind / dm-haiku
JAX-based neural network library
deep-learning deep-neural-networks jax machine-learning neural-networks
Language:Python 2816
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 2798
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 2664
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++ 2424
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++ 2404
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 2367
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 2247
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 2232
google-deepmind / mctx
Monte Carlo tree search in JAX
jax reinforcement-learning monte-carlo-tree-search
Language:Python 2210
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 2179
pennylane
PennyLaneAI / pennylane
PennyLane is a cross-platform Python library for differentiable programming of quantum computers. 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 2129
google / brax
Massively parallel rigidbody physics simulation on accelerator hardware.
jax reinforcement-learning robotics physics-simulation
Language:Jupyter Notebook 2080
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 2057
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 1836
jermwatt / machine_learning_refined
Notes, examples, and Python demos for the 2nd edition of the textbook "Machine Learning Refined" (published by Cambridge University Press).
machine-learning deep-learning artificial-intelligence data-science jupyter-notebook numpy autograd machine-learning-algorithms lecture-notes slides neural-network collab jax mathematical-optimization python
Language:Python 1594
tensorly
tensorly / tensorly
TensorLy: Tensor Learning in Python.
machine-learning tensor decomposition tensor-algebra tensorly tensor-learning python tensor-methods pytorch mxnet jax tensorflow cupy numpy tensor-decomposition tensor-factorization tensor-regression regression
Language:Python 1503
google-research / big_transfer
Official repository for the "Big Transfer (BiT): General Visual Representation Learning" paper.
deep-learning convolutional-neural-networks imagenet tensorflow2 jax pytorch transfer-learning
Language:Python 1494
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 gpu machine-learning python pytorch fp8 jax
Language:Python 1467
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 1465