Projects in Awesome Lists tagged with jax

https://github.com/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

Last synced: 15 Apr 2025

https://github.com/keras-team/keras

Deep Learning for humans

data-science deep-learning jax machine-learning neural-networks python pytorch tensorflow

Last synced: 15 Apr 2025

https://github.com/jax-ml/jax

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more

jax

Last synced: 15 Apr 2025

https://github.com/huggingface/diffusers

🤗 Diffusers: State-of-the-art diffusion models for image, video, and audio generation in PyTorch and FLAX.

deep-learning diffusion flax flux hacktoberfest image-generation image2image image2video jax latent-diffusion-models pytorch score-based-generative-modeling stable-diffusion stable-diffusion-diffusers text2image text2video video2video

Last synced: 15 Apr 2025

https://github.com/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

Last synced: 15 Apr 2025

https://github.com/ml-tooling/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

Last synced: 15 Apr 2025

https://github.com/ivy-llc/ivy

Convert Machine Learning Code Between Frameworks

converter deep-learning ivy jax machine-learning neural-network numpy python pytorch tensorflow translation transpilation

Last synced: 15 Apr 2025

https://github.com/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

Last synced: 15 Apr 2025

https://github.com/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

Last synced: 15 Apr 2025

https://github.com/google/trax

Trax — Deep Learning with Clear Code and Speed

deep-learning deep-reinforcement-learning jax machine-learning numpy reinforcement-learning transformer

Last synced: 08 Apr 2025

https://github.com/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

Last synced: 08 Apr 2025

https://github.com/google/flax

Flax is a neural network library for JAX that is designed for flexibility.

jax

Last synced: 08 Apr 2025

https://github.com/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

Last synced: 13 Apr 2025

https://github.com/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

Last synced: 13 Apr 2025

https://github.com/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

Last synced: 27 Mar 2025

https://github.com/tensorflow/datasets

TFDS is a collection of datasets ready to use with TensorFlow, Jax, ...

data dataset datasets jax machine-learning numpy tensorflow

Last synced: 08 Apr 2025

https://github.com/google-research/scenic

Scenic: A Jax Library for Computer Vision Research and Beyond

attention computer-vision deep-learning jax research transformers vision-transformer

Last synced: 09 Apr 2025

https://github.com/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

Last synced: 13 Apr 2025

https://github.com/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

Last synced: 22 Feb 2025

https://github.com/google-deepmind/dm-haiku

JAX-based neural network library

deep-learning deep-neural-networks jax machine-learning neural-networks

Last synced: 10 Apr 2025

https://github.com/openxla/iree

A retargetable MLIR-based machine learning compiler and runtime toolkit.

compiler cuda jax machine-learning mlir pytorch runtime spirv tensorflow vulkan

Last synced: 09 Dec 2024

https://github.com/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

Last synced: 10 Apr 2025

https://github.com/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

Last synced: 19 Apr 2025

https://github.com/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

Last synced: 08 Apr 2025

https://github.com/google/brax

Massively parallel rigidbody physics simulation on accelerator hardware.

jax physics-simulation reinforcement-learning robotics

Last synced: 08 Apr 2025

https://github.com/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

Last synced: 10 Apr 2025

https://github.com/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

Last synced: 29 Nov 2024

https://github.com/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

Last synced: 14 Mar 2025

https://github.com/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

Last synced: 11 Apr 2025

https://github.com/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.

autograd automatic-differentiation cirq deep-learning differentiable-computing hacktoberfest jax machine-learning neural-network optimization python pytorch qiskit qml quantum quantum-chemistry quantum-computing quantum-machine-learning strawberryfields tensorflow

Last synced: 08 Apr 2025

https://github.com/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.

autograd automatic-differentiation cirq deep-learning differentiable-computing hacktoberfest jax machine-learning neural-network optimization python pytorch qiskit qml quantum quantum-chemistry quantum-computing quantum-machine-learning strawberryfields tensorflow

Last synced: 26 Mar 2025

https://github.com/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

Last synced: 11 Apr 2025

https://github.com/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

Last synced: 11 Apr 2025

https://github.com/google-deepmind/mctx

Monte Carlo tree search in JAX

jax monte-carlo-tree-search reinforcement-learning

Last synced: 27 Mar 2025

https://github.com/deepmind/mctx

Monte Carlo tree search in JAX

jax monte-carlo-tree-search reinforcement-learning

Last synced: 21 Dec 2024

https://github.com/patrick-kidger/equinox

Elegant easy-to-use neural networks + scientific computing in JAX. https://docs.kidger.site/equinox/

deep-learning equinox jax neural-networks

Last synced: 08 Apr 2025

https://github.com/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

Last synced: 10 Apr 2025

https://github.com/apple/axlearn

An Extensible Deep Learning Library

deep-learning jax

Last synced: 11 Apr 2025

https://github.com/google-deepmind/penzai

A JAX research toolkit for building, editing, and visualizing neural networks.

fine-tuning interpretability jax neural-networks visualization

Last synced: 10 Apr 2025

https://github.com/kubeflow/trainer

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

Last synced: 10 Apr 2025

https://github.com/neonwatty/machine_learning_refined

Notes, Python demos / notebooks, and free chapters for the 2nd edition of the university textbook "Machine Learning Refined".

artificial-intelligence autograd collab data-science deep-learning genai jax jupyter-notebook lecture-notes machine-learning machine-learning-algorithms mathematical-optimization neural-network numpy python slides

Last synced: 12 Dec 2024

https://github.com/deepmodeling/deepmd-kit

A deep learning package for many-body potential energy representation and molecular dynamics

ase c computational-chemistry cpp cuda deep-learning deepmd ipi jax lammps materials-science molecular-dynamics nodejs paddle potential-energy python pytorch rocm tensorflow

Last synced: 09 Apr 2025

https://github.com/danijar/dreamerv3

Mastering Diverse Domains through World Models

artificial-intelligence general jax minecraft reinforcement-learning world-models

Last synced: 10 Apr 2025

https://github.com/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

Last synced: 06 Feb 2025

https://github.com/patrick-kidger/diffrax

Numerical differential equation solvers in JAX. Autodifferentiable and GPU-capable. https://docs.kidger.site/diffrax/

deep-learning differential-equations dynamical-systems equinox jax machine-learning neural-differential-equations neural-networks

Last synced: 09 Apr 2025

https://github.com/kubeflow/katib

Automated Machine Learning on Kubernetes

ai automl huggingface hyperparameter-tuning jax kubeflow kubernetes llm machine-learning mlops neural-architecture-search pytorch scikit-learn tensorflow

Last synced: 09 Apr 2025

https://github.com/tensorly/tensorly

TensorLy: Tensor Learning in Python.

cupy decomposition jax machine-learning mxnet numpy python pytorch regression tensor tensor-algebra tensor-decomposition tensor-factorization tensor-learning tensor-methods tensor-regression tensorflow tensorly

Last synced: 09 Apr 2025

https://github.com/google-research/big_transfer

Official repository for the "Big Transfer (BiT): General Visual Representation Learning" paper.

convolutional-neural-networks deep-learning imagenet jax pytorch tensorflow2 transfer-learning

Last synced: 07 Apr 2025

https://github.com/yang-song/score_sde

Official code for Score-Based Generative Modeling through Stochastic Differential Equations (ICLR 2021, Oral)

controllable-generation diffusion-models flax generative-models iclr-2021 inverse-problems jax score-based-generative-modeling score-matching stochastic-differential-equations

Last synced: 08 Apr 2025

https://github.com/patrick-kidger/diffrax?ref=zzun.app

Numerical differential equation solvers in JAX. Autodifferentiable and GPU-capable. https://docs.kidger.site/diffrax/

deep-learning differential-equations dynamical-systems equinox jax machine-learning neural-differential-equations neural-networks

Last synced: 13 Nov 2024

https://github.com/deepmind/jraph

A Graph Neural Network Library in Jax

deep-learning graph-neural-networks jax machine-learning

Last synced: 31 Mar 2025

https://github.com/google-deepmind/jraph

A Graph Neural Network Library in Jax

deep-learning graph-neural-networks jax machine-learning

Last synced: 08 Apr 2025

https://github.com/swanhubx/swanlab

⚡️SwanLab - an open-source, modern-design AI training tracking and visualization tool. Supports Cloud / Self-hosted use. Integrated with PyTorch / Transformers / LLaMA Factory / Swift / Ultralytics / veRL / MMEngine / Keras etc.

data-science deep-learning jax logging machine-learning mlops model-versioning python pytorch tensorboard tensorflow tracking transformers visualization

Last synced: 13 Apr 2025

https://github.com/apple/ml-aim

This repository provides the code and model checkpoints for AIMv1 and AIMv2 research projects.

jax large-scale-vision-models mlx pytorch

Last synced: 10 Apr 2025

https://github.com/patrick-kidger/jaxtyping

Type annotations and runtime checking for shape and dtype of JAX/NumPy/PyTorch/etc. arrays. https://docs.kidger.site/jaxtyping/

jax python-typing typing

Last synced: 13 Apr 2025

https://github.com/dbraun/dawdreamer

Digital Audio Workstation with Python; VST instruments/effects, parameter automation, FAUST, JAX, Warp Markers, and JUCE processors

ableton audio audio-plugin audio-processing daw faust jax juce midi python synthesizer vst vst-host vst3 vst3-host

Last synced: 10 Apr 2025

https://github.com/DBraun/DawDreamer

Digital Audio Workstation with Python; VST instruments/effects, parameter automation, FAUST, JAX, Warp Markers, and JUCE processors

ableton audio audio-plugin audio-processing daw faust jax juce midi python synthesizer vst vst-host vst3 vst3-host

Last synced: 16 Mar 2025

https://github.com/google/jaxopt

Hardware accelerated, batchable and differentiable optimizers in JAX.

bi-level deep-learning differentiable-programming jax optimization

Last synced: 27 Mar 2025

https://github.com/awslabs/fortuna

A Library for Uncertainty Quantification.

ai bayesian-inference calibration conformal-prediction deep-learning flax jax machine-learning ml model-calibration neural-networks uncertainty uncertainty-calibration uncertainty-estimation uncertainty-quantification

Last synced: 27 Mar 2025

https://github.com/keras-team/keras-hub

Pretrained model hub for Keras 3.

deep-learning jax keras llm machine-learning natural-language-processing nlp python pytorch tensorflow

Last synced: 10 Apr 2025

https://github.com/pdebench/pdebench

PDEBench: An Extensive Benchmark for Scientific Machine Learning

ai autoregressive-models benchmark deep-learning fluid-dynamics jax machine-learning navier-stokes-equations neural-networks neural-operators partial-differential-equations physics-informed-neural-networks pytorch scientific scientific-computing sciml simulation

Last synced: 02 Apr 2025

https://github.com/mosecorg/mosec

A high-performance ML model serving framework, offers dynamic batching and CPU/GPU pipelines to fully exploit your compute machine

cv deep-learning gpu hacktoberfest jax llm llm-serving machine-learning machine-learning-platform mlops model-serving mxnet nerual-network python pytorch rust tensorflow tts

Last synced: 12 Apr 2025

https://github.com/parrt/tensor-sensor

The goal of this library is to generate more helpful exception messages for matrix algebra expressions for numpy, pytorch, jax, tensorflow, keras, fastai.

debugging deep-learning jax matrix numpy python pytorch tensorflow tracing vector

Last synced: 08 Apr 2025

https://github.com/instadeepai/mava

🦁 A research-friendly codebase for fast experimentation of multi-agent reinforcement learning in JAX

jax marl multi-agent-reinforcement-learning multi-agent-systems multiagent reinforcement-learning research

Last synced: 12 Apr 2025

https://github.com/google/jax-cfd

Computational Fluid Dynamics in JAX

cfd jax

Last synced: 14 Mar 2025

https://github.com/google-research/long-range-arena

Long Range Arena for Benchmarking Efficient Transformers

attention deep-learning flax jax nlp transformers

Last synced: 22 Feb 2025

https://github.com/probml/dynamax

State Space Models library in JAX

hidden-markov-models jax kalman-filter python state-space-models

Last synced: 13 Apr 2025

https://github.com/instadeepai/Mava

🦁 A research-friendly codebase for fast experimentation of multi-agent reinforcement learning in JAX

jax marl multi-agent-reinforcement-learning multi-agent-systems multiagent reinforcement-learning research

Last synced: 13 Apr 2025

https://github.com/instadeepai/jumanji

🕹️ A diverse suite of scalable reinforcement learning environments in JAX

jax python reinforcement-learning research

Last synced: 11 Apr 2025

https://github.com/gordicaleksa/get-started-with-jax

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.

deep-learning flax haiku jax jupyter lax learn-jax machine-learning numpy optax python tutorial xla

Last synced: 04 Apr 2025

https://github.com/jonasrauber/eagerpy

PyTorch, TensorFlow, JAX and NumPy — all of them natively using the same code

eager-execution jax numpy python pytorch tensorflow tensorflow2

Last synced: 08 Apr 2025

https://github.com/google-research/prompt-tuning

Original Implementation of Prompt Tuning from Lester, et al, 2021

flax jax language-model machine-learning nlp prompt-tuning

Last synced: 10 Apr 2025

https://github.com/luchris429/purejaxrl

Really Fast End-to-End Jax RL Implementations

deep-reinforcement-learning jax ppo reinforcement-learning reinforcement-learning-algorithms

Last synced: 20 Mar 2025

https://github.com/mikeroyal/machine-learning-guide

Machine learning Guide. Learn all about Machine Learning Tools, Libraries, Frameworks, Large Language Models (LLMs), and Training Models.

artificial-neural-networks aws-sagemaker deep-learning generative-ai gpt-3 gpt-4 gpt4all image-classification image-processing jax llms machine-learning-library machine-learning-models machinelearning machinelearning-python python pytorch scikit-learn scikitlearn-machine-learning support-vector-machines

Last synced: 05 Apr 2025

https://github.com/ikostrikov/jaxrl

JAX (Flax) implementation of algorithms for Deep Reinforcement Learning with continuous action spaces.

batch-reinforcement-learning behavioral-cloning continuous-control deep-deterministic-policy-gradient deep-learning deep-reinforcement-learning flax gym jax offline-reinforcement-learning reinforcement-learning sac soft-actor-critic

Last synced: 04 Apr 2025

https://github.com/google-research/fast-soft-sort

Fast Differentiable Sorting and Ranking

differentiable jax pytorch ranking sorting tensorflow

Last synced: 04 Apr 2025

https://github.com/facebookresearch/neuralcompression

A collection of tools for neural compression enthusiasts.

compression deep-learning jax machine-learning neural-compression python pytorch

Last synced: 06 Apr 2025

https://github.com/netket/netket

Machine learning algorithms for many-body quantum systems

complex-neural-network deep-learning exact-diagonalization hamiltonian jax machine-learning machine-learning-algorithms markov-chain-monte-carlo monte-carlo-methods neural-networks physics-simulation quantum quantum-state-tomography unitaryhack variational-method variational-monte-carlo

Last synced: 02 Apr 2025

https://github.com/Toni-SM/skrl

Modular reinforcement learning library (on PyTorch and JAX) with support for NVIDIA Isaac Gym, Omniverse Isaac Gym and Isaac Lab

deep-learning deepmind gym gymnasium isaac-gym isaac-lab isaac-orbit isaac-sim isaaclab jax machine-learning nvidia-omniverse openai-gym python pytorch reinforcement-learning rl robosuite robotics skrl

Last synced: 02 Apr 2025

https://github.com/ayaka14732/tpu-starter

Everything you want to know about Google Cloud TPU

cloud-tpu deep-learning gcp google-cloud-platform jax machine-learning tpu

Last synced: 04 Apr 2025

https://github.com/facebookresearch/NeuralCompression

A collection of tools for neural compression enthusiasts.

compression deep-learning jax machine-learning neural-compression python pytorch

Last synced: 23 Nov 2024

https://github.com/ott-jax/ott

Optimal transport tools implemented with the JAX framework, to get differentiable, parallel and jit-able computations.

automatic-differentiation gromov-wasserstein jax optimal-transport sinkhorn

Last synced: 12 Nov 2024

https://github.com/SwanHubX/SwanLab

⚡️SwanLab: your ML experiment notebook. 你的AI实验笔记本，日志记录与可视化AI训练全流程。

data-science deep-learning fastapi jax machine-learning mlops model-versioning python pytorch tensorboard tensorflow tracking transformers visualization

Last synced: 05 Mar 2025

https://github.com/srush/annotated-s4

Implementation of https://srush.github.io/annotated-s4

deep-learning jax

Last synced: 05 Apr 2025

https://srush.github.io/annotated-s4/

Implementation of https://srush.github.io/annotated-s4

deep-learning jax

Last synced: 20 Mar 2025

https://github.com/samuela/git-re-basin

Code release for "Git Re-Basin: Merging Models modulo Permutation Symmetries"

deep-learning deeplearning jax machine-learning neural-networks

Last synced: 05 Apr 2025

https://github.com/poets-ai/elegy

A High Level API for Deep Learning in JAX

deep-learning jax

Last synced: 12 Apr 2025

https://poets-ai.github.io/elegy/

A High Level API for Deep Learning in JAX

deep-learning jax

Last synced: 18 Nov 2024

https://github.com/google/paxml

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.

c4 gpt jax large-language-models llm model-flops parallelism