An open API service providing repository metadata for many open source software ecosystems.

GitHub topics: jax

kach/memo

A domain-specific probabilistic programming language for reasoning about reasoning

Language: Python - Size: 42.1 MB - Last synced at: about 3 hours ago - Pushed at: about 4 hours ago - Stars: 68 - Forks: 4

mfschubert/fmmax

Fourier Modal Method with Jax

Language: Python - Size: 1.45 MB - Last synced at: about 5 hours ago - Pushed at: about 6 hours ago - Stars: 5 - Forks: 1

tensorcircuit/tensorcircuit-ng

Tensor network based quantum software framework: next generation

Language: Python - Size: 15.3 MB - Last synced at: about 6 hours ago - Pushed at: about 6 hours ago - Stars: 19 - Forks: 3

synsense/rockpool

A machine learning library for spiking neural networks. Supports training with both torch and jax pipelines, and deployment to neuromorphic hardware.

Language: Python - Size: 830 MB - Last synced at: about 7 hours ago - Pushed at: about 8 hours ago - Stars: 65 - Forks: 13

tumaer/JAXFLUIDS

Differentiable Fluid Dynamics Package

Language: Python - Size: 12.5 MB - Last synced at: about 8 hours ago - Pushed at: about 9 hours ago - Stars: 402 - Forks: 72

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.

Language: Python - Size: 117 MB - Last synced at: about 9 hours ago - Pushed at: about 12 hours ago - Stars: 2,620 - Forks: 657

google-research/kauldron

Modular, scalable library to train ML models

Language: Python - Size: 1.79 MB - Last synced at: about 10 hours ago - Pushed at: about 11 hours ago - Stars: 111 - Forks: 13

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.

Language: Python - Size: 303 MB - Last synced at: about 21 hours ago - Pushed at: 9 months ago - Stars: 25,764 - Forks: 4,619

iree-org/iree

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

Language: C++ - Size: 211 MB - Last synced at: about 21 hours ago - Pushed at: about 21 hours ago - Stars: 3,121 - Forks: 694

lululxvi/deepxde

A library for scientific machine learning and physics-informed learning

Language: Python - Size: 27.2 MB - Last synced at: about 22 hours ago - Pushed at: about 24 hours ago - Stars: 3,188 - Forks: 819

patrick-kidger/equinox

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

Language: Python - Size: 39.8 MB - Last synced at: 1 day ago - Pushed at: 1 day ago - Stars: 2,345 - Forks: 161

pmocz/advectiondiffusion-jax

Solve the advection diffusion equations looped into an optimization problem with JAX/autodiff

Language: Python - Size: 371 KB - Last synced at: 1 day ago - Pushed at: 1 day ago - Stars: 11 - Forks: 3

acrovato/pybench

Benchmark of Python codes

Language: Python - Size: 13.7 KB - Last synced at: 1 day ago - Pushed at: 1 day ago - Stars: 0 - Forks: 0

cabouman/mbirjax

MBIRJAX is a Python package for Model Based Iterative Reconstruction (MBIR) of images from tomographic data.

Language: Python - Size: 2.55 MB - Last synced at: 1 day ago - Pushed at: 1 day ago - Stars: 13 - Forks: 6

jeree02847/Dream

Dream is a collaborative platform that helps users explore and share their creative visions. It offers tools for brainstorming, visual storytelling, and community feedback, making idea development easier and more engaging.

Size: 3.91 KB - Last synced at: 1 day ago - Pushed at: 1 day ago - Stars: 0 - Forks: 0

nobodyPerfecZ/vit-finetuned-patch16-224-recaptchav2-v1

An AI-powered computer vision system written in PyTorch, Tensorflow and JAX, designed to analyze and solve reCAPTCHA v2 challenges.

Language: Jupyter Notebook - Size: 7.81 KB - Last synced at: 1 day ago - Pushed at: 1 day ago - Stars: 0 - Forks: 0

arish-mhrjn/aimodelinspector

A fairly comprehensive Python library allowing for exploration, self-education and categorizaton of AI models

Language: Python - Size: 214 KB - Last synced at: 1 day ago - Pushed at: 1 day ago - Stars: 1 - Forks: 0

explosion/thinc

🔮 A refreshing functional take on deep learning, compatible with your favorite libraries

Language: Python - Size: 10.6 MB - Last synced at: about 15 hours ago - Pushed at: about 1 month ago - Stars: 2,848 - Forks: 280

debangshu-mukherjee/ptyrodactyl

JAX based ptychography reconstructions through differentiable programming.

Language: Jupyter Notebook - Size: 1.33 MB - Last synced at: 2 days ago - Pushed at: 2 days ago - Stars: 0 - Forks: 0

JoeJoe1313/PaliGemma-Image-Segmentation

An app with FastAPI, Docker, transformers, JAX/Flax for performing image segmentation with PaliGemma 2 mix

Language: Python - Size: 7.46 MB - Last synced at: 2 days ago - Pushed at: 2 days ago - Stars: 0 - Forks: 0

flaport/klujax

Solve sparse linear systems in JAX using the KLU algorithm

Language: Python - Size: 410 KB - Last synced at: 2 days ago - Pushed at: 2 days ago - Stars: 35 - Forks: 4

robochimps/vibrojet

Python package for constructing molecular rovibrational kinetic and potential energy operators using Taylor-mode automatic differentiation

Language: Python - Size: 5.5 MB - Last synced at: 2 days ago - Pushed at: 2 days ago - Stars: 0 - Forks: 0

lucidrains/flash-attention-jax

Implementation of Flash Attention in Jax

Language: Python - Size: 181 KB - Last synced at: 2 days ago - Pushed at: about 1 year ago - Stars: 208 - Forks: 24

AI-Hypercomputer/JetStream

JetStream is a throughput and memory optimized engine for LLM inference on XLA devices, starting with TPUs (and GPUs in future -- PRs welcome).

Language: Python - Size: 6.35 MB - Last synced at: about 2 hours ago - Pushed at: about 3 hours ago - Stars: 322 - Forks: 39

n2cholas/awesome-jax

JAX - A curated list of resources https://github.com/google/jax

Size: 353 KB - Last synced at: 1 day ago - Pushed at: 3 months ago - Stars: 1,802 - Forks: 143

pasqal-io/qadence

Digital-analog quantum programming interface

Language: Python - Size: 51.7 MB - Last synced at: about 15 hours ago - Pushed at: 2 days ago - Stars: 85 - Forks: 22

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.

Language: Python - Size: 41.7 MB - Last synced at: 2 days ago - Pushed at: 2 days ago - Stars: 1,470 - Forks: 96

pdebench/PDEBench

PDEBench: An Extensive Benchmark for Scientific Machine Learning

Language: Python - Size: 1.09 MB - Last synced at: 2 days ago - Pushed at: 3 months ago - Stars: 891 - Forks: 104

danielward27/flowjax

Language: Python - Size: 18 MB - Last synced at: 2 days ago - Pushed at: 2 days ago - Stars: 158 - Forks: 16

patrick-kidger/lineax

Linear solvers in JAX and Equinox. https://docs.kidger.site/lineax

Language: Python - Size: 192 KB - Last synced at: 2 days ago - Pushed at: 2 days ago - Stars: 439 - Forks: 27

ami-iit/jaxsim

A differentiable physics engine and multibody dynamics library for control and robot learning.

Language: Python - Size: 3.15 MB - Last synced at: 2 days ago - Pushed at: 3 days ago - Stars: 118 - Forks: 14

EMI-Group/evox

Distributed GPU-Accelerated Framework for Evolutionary Computation. Comprehensive Library of Evolutionary Algorithms & Benchmark Problems.

Language: Python - Size: 42.5 MB - Last synced at: 2 days ago - Pushed at: 3 days ago - Stars: 498 - Forks: 75

bahremsd/tmmax-workshop

Workshop given in graduate-level thin film coatings course in ITU

Language: Jupyter Notebook - Size: 2.73 MB - Last synced at: 3 days ago - Pushed at: 3 days ago - Stars: 0 - Forks: 0

QuantEcon/lecture-python.myst

Quantitative Economics with Python

Language: Jupyter Notebook - Size: 3.12 GB - Last synced at: 3 days ago - Pushed at: 3 days ago - Stars: 97 - Forks: 45

JaxGaussianProcesses/GPJax

Gaussian processes in JAX and Flax.

Language: Python - Size: 279 MB - Last synced at: 1 day ago - Pushed at: 26 days ago - Stars: 507 - Forks: 63

froystig/jax Fork of jax-ml/jax

my jax fork

Language: Python - Size: 110 MB - Last synced at: 3 days ago - Pushed at: 3 days ago - Stars: 1 - Forks: 0

HMUNACHI/nanodl

A Jax-based library for building transformers, includes implementations of GPT, Gemma, LlaMa, Mixtral, Whisper, SWin, ViT and more.

Language: Python - Size: 44.4 MB - Last synced at: 1 day ago - Pushed at: 8 months ago - Stars: 287 - Forks: 10

NACLab/ngc-learn

NGC-Learn: Neurobiological Systems Design and Simulation in Python

Language: Python - Size: 79.9 MB - Last synced at: 3 days ago - Pushed at: 3 days ago - Stars: 148 - Forks: 30

invrs-io/opt

Algorithms for inverse design

Language: Python - Size: 182 KB - Last synced at: 3 days ago - Pushed at: 3 days ago - Stars: 7 - Forks: 1

google/grain

Library for reading and processing ML training data.

Language: Python - Size: 4.29 MB - Last synced at: 3 days ago - Pushed at: 3 days ago - Stars: 434 - Forks: 35

keras-team/keras-hub

Pretrained model hub for Keras 3.

Language: Python - Size: 6.72 MB - Last synced at: 3 days ago - Pushed at: 3 days ago - Stars: 890 - Forks: 278

sotetsuk/pgx

♟️ Vectorized RL game environments in JAX

Language: Python - Size: 74.6 MB - Last synced at: about 8 hours ago - Pushed at: 2 months ago - Stars: 473 - Forks: 34

sanchit-gandhi/whisper-jax

JAX implementation of OpenAI's Whisper model for up to 70x speed-up on TPU.

Language: Jupyter Notebook - Size: 8.75 MB - Last synced at: 3 days ago - Pushed at: about 1 year ago - Stars: 4,590 - Forks: 402

google/jaxonnxruntime

A user-friendly tool chain that enables the seamless execution of ONNX models using JAX as the backend.

Language: Python - Size: 425 KB - Last synced at: 1 day ago - Pushed at: 1 day ago - Stars: 111 - Forks: 17

google/orbax

Orbax provides common checkpointing and persistence utilities for JAX users

Language: Python - Size: 4.34 MB - Last synced at: 3 days ago - Pushed at: 3 days ago - Stars: 376 - Forks: 48

KangOxford/AlphaTrade

JAX-LOB: A GPU-Accelerated limit order book simulator to unlock large scale reinforcement learning for trading

Language: Jupyter Notebook - Size: 240 MB - Last synced at: 3 days ago - Pushed at: 3 days ago - Stars: 121 - Forks: 24

ml-tooling/best-of-ml-python

🏆 A ranked list of awesome machine learning Python libraries. Updated weekly.

Size: 19.1 MB - Last synced at: 3 days ago - Pushed at: 8 days ago - Stars: 20,017 - Forks: 2,766

ComputationalPsychiatry/pyhgf

PyHGF: A neural network library for predictive coding

Language: Python - Size: 961 MB - Last synced at: 3 days ago - Pushed at: 3 days ago - Stars: 74 - Forks: 21

deepmodeling/deepmd-kit

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

Language: Python - Size: 63.5 MB - Last synced at: 3 days ago - Pushed at: 4 days ago - Stars: 1,635 - Forks: 544

sxzgroup/quantum_excited_state

Code implementation and data for arXiv:2504.21459

Language: Python - Size: 73.5 MB - Last synced at: 4 days ago - Pushed at: 4 days ago - Stars: 1 - Forks: 0

google/trax

Trax — Deep Learning with Clear Code and Speed

Language: Python - Size: 162 MB - Last synced at: 3 days ago - Pushed at: 29 days ago - Stars: 8,203 - Forks: 825

texttron/tevatron

Tevatron - Unified Document Retrieval Toolkit across Scale, Language, and Modality. Demo in SIGIR 2023, SIGIR 2025.

Language: Python - Size: 20.4 MB - Last synced at: about 22 hours ago - Pushed at: about 22 hours ago - Stars: 592 - Forks: 104

google-deepmind/dm-haiku

JAX-based neural network library

Language: Python - Size: 2.5 MB - Last synced at: 3 days ago - Pushed at: 8 days ago - Stars: 3,020 - Forks: 244

wcxve/xspex

Access Xspec models and corresponding JAX/XLA ops.

Language: Python - Size: 67.4 KB - Last synced at: 4 days ago - Pushed at: 4 days ago - Stars: 4 - Forks: 1

liblaf/apple

🍎 Advanced optimization and math utilities with JAX for Python projects!

Language: Python - Size: 2.31 MB - Last synced at: 3 days ago - Pushed at: 3 days ago - Stars: 0 - Forks: 0

nv-legate/multimesh-jax

PjRt plugin and Python APIs for MPMD workflows in Jax

Language: HTML - Size: 27.8 MB - Last synced at: 4 days ago - Pushed at: 4 days ago - Stars: 0 - Forks: 0

tensorly/tensorly

TensorLy: Tensor Learning in Python.

Language: Python - Size: 16.1 MB - Last synced at: 4 days ago - Pushed at: 4 days ago - Stars: 1,608 - Forks: 293

lenstronomy/JAXtronomy

JAX port of lenstronomy, for parallelized, GPU accelerated, and differentiable gravitational lensing and image simulations.

Language: Python - Size: 1.12 MB - Last synced at: 4 days ago - Pushed at: 4 days ago - Stars: 17 - Forks: 6

erfanzar/EasyDeL

Accelerate, Optimize performance with streamlined training and serving options with JAX.

Language: Python - Size: 15 MB - Last synced at: 4 days ago - Pushed at: 4 days ago - Stars: 268 - Forks: 28

erfanzar/eformer

(EasyDel Former) is a utility library designed to simplify and enhance the development in JAX

Language: Python - Size: 3.2 MB - Last synced at: 3 days ago - Pushed at: 4 days ago - Stars: 27 - Forks: 3

PennyLaneAI/catalyst

A JIT compiler for hybrid quantum programs in PennyLane

Language: Python - Size: 14.9 MB - Last synced at: 4 days ago - Pushed at: 4 days ago - Stars: 162 - Forks: 46

WICKED4950/BwETAF-IID-100M

New more optimized 100M LLM trained from scratch using JAX

Size: 4.88 KB - Last synced at: 4 days ago - Pushed at: 4 days ago - Stars: 0 - Forks: 0

google-research/scenic

Scenic: A Jax Library for Computer Vision Research and Beyond

Language: Python - Size: 63.7 MB - Last synced at: 4 days ago - Pushed at: 4 days ago - Stars: 3,529 - Forks: 454

huggingface/diffusers

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

Language: Python - Size: 65.8 MB - Last synced at: 4 days ago - Pushed at: 4 days ago - Stars: 28,849 - Forks: 5,924

huggingface/transformers

🤗 Transformers: State-of-the-art Machine Learning for Pytorch, TensorFlow, and JAX.

Language: Python - Size: 288 MB - Last synced at: 4 days ago - Pushed at: 4 days ago - Stars: 143,865 - Forks: 28,851

jax-ml/jax

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

Language: Python - Size: 114 MB - Last synced at: 4 days ago - Pushed at: 4 days ago - Stars: 32,125 - Forks: 3,010

team-ocean/veros

The versatile ocean simulator, in pure Python, powered by JAX.

Language: Python - Size: 28.3 MB - Last synced at: 3 days ago - Pushed at: 24 days ago - Stars: 347 - Forks: 56

keras-team/keras

Deep Learning for humans

Language: Python - Size: 44.3 MB - Last synced at: 4 days ago - Pushed at: 5 days ago - Stars: 62,949 - Forks: 19,576

mlaux1/rl-blox

Modular JAX-based toolbox for implementing RL algorithms.

Language: Python - Size: 1.25 MB - Last synced at: 3 days ago - Pushed at: 3 days ago - Stars: 7 - Forks: 1

iit-DLSLab/Quadruped-PyMPC

A model predictive controller for quadruped robots based on the single rigid body model and written in python. Gradient-based (acados) or Sampling-based (jax)

Language: Python - Size: 44.4 MB - Last synced at: 5 days ago - Pushed at: 5 days ago - Stars: 297 - Forks: 24

pasteurlabs/tesseract-jax

Execute + differentiate Tesseracts as part of JAX programs, with full support for function transformations like JIT, grad, and more. ⚡

Language: Python - Size: 5.24 MB - Last synced at: 5 days ago - Pushed at: 5 days ago - Stars: 14 - Forks: 0

wandb/wandb

The AI developer platform. Use Weights & Biases to train and fine-tune models, and manage models from experimentation to production.

Language: Python - Size: 171 MB - Last synced at: 4 days ago - Pushed at: 5 days ago - Stars: 9,832 - Forks: 736

deepmodeling/jax-fem

Differentiable Finite Element Method with JAX

Language: Python - Size: 79.1 MB - Last synced at: 5 days ago - Pushed at: 5 days ago - Stars: 383 - Forks: 62

arogozhnikov/einops

Flexible and powerful tensor operations for readable and reliable code (for pytorch, jax, TF and others)

Language: Python - Size: 2.6 MB - Last synced at: 5 days ago - Pushed at: 15 days ago - Stars: 8,885 - Forks: 369

felafax/felafax

Felafax is building AI infra for non-NVIDIA GPUs

Language: Jupyter Notebook - Size: 3.36 MB - Last synced at: 4 days ago - Pushed at: 4 months ago - Stars: 559 - Forks: 35

jaco-bro/nnx-lm

nnx-lm: A portable, pip-installable CLI for running LLMs via JAX on any hardware backend.

Language: Python - Size: 0 Bytes - Last synced at: 5 days ago - Pushed at: 5 days ago - Stars: 0 - Forks: 0

ashishpatel26/Treasure-of-Transformers

💁 Awesome Treasure of Transformers Models for Natural Language processing contains papers, videos, blogs, official repo along with colab Notebooks. 🛫☑️

Language: Jupyter Notebook - Size: 370 KB - Last synced at: 4 days ago - Pushed at: 10 months ago - Stars: 987 - Forks: 210

benmoseley/FBPINNs

Solve forward and inverse problems related to partial differential equations using finite basis physics-informed neural networks (FBPINNs)

Language: Python - Size: 17.2 MB - Last synced at: 5 days ago - Pushed at: 5 days ago - Stars: 390 - Forks: 78

google/flax

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

Language: Jupyter Notebook - Size: 24.9 MB - Last synced at: 4 days ago - Pushed at: 6 days ago - Stars: 6,527 - Forks: 695

ivy-llc/ivy

Convert Machine Learning Code Between Frameworks

Language: Python - Size: 166 MB - Last synced at: 4 days ago - Pushed at: 11 days ago - Stars: 14,197 - Forks: 5,665

grame-cncm/faust

Functional programming language for signal processing and sound synthesis

Language: C++ - Size: 362 MB - Last synced at: 6 days ago - Pushed at: 6 days ago - Stars: 2,719 - Forks: 347

dc-setyadi/nuni-neomu-areumdawo

My profile readme!!!

Size: 8.79 KB - Last synced at: 6 days ago - Pushed at: 6 days ago - Stars: 0 - Forks: 0

HomebrewML/Olmax

HomebrewNLP in JAX flavour for maintable TPU-Training

Language: Python - Size: 6.03 MB - Last synced at: 3 days ago - Pushed at: over 1 year ago - Stars: 50 - Forks: 6

ott-jax/ott

Optimal transport tools implemented with the JAX framework, to solve large scale matching problems of any flavor.

Language: Python - Size: 164 MB - Last synced at: 6 days ago - Pushed at: 6 days ago - Stars: 586 - Forks: 114

apple/axlearn

An Extensible Deep Learning Library

Language: Python - Size: 21.1 MB - Last synced at: 6 days ago - Pushed at: 6 days ago - Stars: 2,031 - Forks: 316

scbirlab/duvida

🧐 Calculating exact and approximate confidence and information metrics for deep learning on general purpose and chemistry tasks.

Language: Python - Size: 21.6 MB - Last synced at: 3 days ago - Pushed at: 3 days ago - Stars: 0 - Forks: 0

sashakolpakov/dire-jax

DImensionality REduction in JAX

Language: Python - Size: 111 MB - Last synced at: about 7 hours ago - Pushed at: about 1 month ago - Stars: 4 - Forks: 1

mashaan14/YouTube-channel

Code I used for my YouTube videos

Language: Jupyter Notebook - Size: 1.38 GB - Last synced at: 6 days ago - Pushed at: 6 days ago - Stars: 6 - Forks: 2

bminixhofer/tokenkit

A toolkit implementing advanced methods to transfer models and model knowledge across tokenizers.

Language: Python - Size: 463 KB - Last synced at: 4 days ago - Pushed at: 4 days ago - Stars: 18 - Forks: 2

mila-iqia/ResearchTemplate

Research Project Template Repository

Language: Python - Size: 2.19 MB - Last synced at: 5 days ago - Pushed at: 14 days ago - Stars: 29 - Forks: 6

SoDevLog/PyTrading

Discover the trading in Python with many strategies, including deep learning approaches. TradingInPython is an open software platform that empowers traders with advanced tools for stock market analysis and trading strategies. Leverage data-driven insights and automate your trading decisions with an open, flexible solution.

Language: Tcl - Size: 284 MB - Last synced at: 7 days ago - Pushed at: 7 days ago - Stars: 2 - Forks: 0

patrick-kidger/quax

Multiple dispatch over abstract array types in JAX.

Language: Python - Size: 144 KB - Last synced at: 7 days ago - Pushed at: 7 days ago - Stars: 120 - Forks: 4

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.

Language: Python - Size: 9.58 MB - Last synced at: 7 days ago - Pushed at: 7 days ago - Stars: 2,390 - Forks: 416

brentyi/jaxls

Sparse nonlinear least squares in JAX

Language: Python - Size: 18.2 MB - Last synced at: 7 days ago - Pushed at: 7 days ago - Stars: 195 - Forks: 13

SciML/SciMLBenchmarks.jl

Scientific machine learning (SciML) benchmarks, AI for science, and (differential) equation solvers. Covers Julia, Python (PyTorch, Jax), MATLAB, R

Language: MATLAB - Size: 186 MB - Last synced at: 7 days ago - Pushed at: 7 days ago - Stars: 328 - Forks: 92

SciML/SciMLBenchmarksOutput

SciML-Bench Benchmarks for Scientific Machine Learning (SciML), Physics-Informed Machine Learning (PIML), and Scientific AI Performance

Language: HTML - Size: 856 MB - Last synced at: 1 day ago - Pushed at: 7 days ago - Stars: 21 - Forks: 5

google-deepmind/penzai

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

Language: Python - Size: 484 MB - Last synced at: 6 days ago - Pushed at: 14 days ago - Stars: 1,770 - Forks: 63

kubeflow/trainer

Distributed ML Training and Fine-Tuning on Kubernetes

Language: Python - Size: 100 MB - Last synced at: 6 days ago - Pushed at: 6 days ago - Stars: 1,773 - Forks: 775

ndarwin314/ninjax

(hopefully) fast pokemon simulator written in JAX to be run on GPU

Language: Python - Size: 249 KB - Last synced at: 7 days ago - Pushed at: 7 days ago - Stars: 1 - Forks: 0

apple/ml-aim

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

Language: Python - Size: 797 KB - Last synced at: 6 days ago - Pushed at: 17 days ago - Stars: 1,275 - Forks: 59

Related Keywords
jax 1,201 deep-learning 256 machine-learning 224 python 218 flax 176 pytorch 170 tensorflow 116 reinforcement-learning 109 numpy 64 neural-networks 53 neural-network 44 transformer 44 transformers 41 keras 41 automatic-differentiation 40 optimization 40 optax 40 ai 37 nlp 33 computer-vision 33 haiku 31 gpu 30 llm 28 equinox 27 tpu 23 bayesian-inference 23 deep-reinforcement-learning 23 dm-haiku 21 differentiable-programming 21 diffusion-models 21 torch 20 python3 20 artificial-intelligence 18 physics 18 simulation 17 probabilistic-programming 17 autodiff 16 cuda 16 xla 16 ml 16 huggingface 16 autograd 15 numpyro 14 gaussian-processes 14 gpt 14 natural-language-processing 13 research 13 large-language-models 13 data-science 13 ppo 13 graph-neural-networks 13 astrophysics 13 scientific-computing 12 tensorflow2 12 astronomy 12 reinforcement-learning-algorithms 11 quantum-computing 11 rl 11 variational-inference 11 mxnet 11 gym 11 deep-neural-networks 11 vision-transformer 11 attention 11 keras3 11 language-model 11 vae 10 brax 10 mujoco 10 robotics 10 tutorial 10 diffusion 10 gpu-acceleration 10 quantum-chemistry 10 jit 9 scipy 9 inverse-problems 9 optics 9 scikit-learn 9 pde 9 mlops 9 quantum 9 sac 9 physics-informed-neural-networks 9 generative-model 9 visualization 9 pytorch-lightning 8 meta-learning 8 cnn 8 optimization-algorithms 8 neuroevolution 8 dynamical-systems 8 physics-simulation 8 score-based-generative-modeling 8 hpc 8 differential-equations 8 rnn 8 td3 8 dqn 8 cupy 8