Awesome JAX

JAXを扱う資料や関連プロジェクトをまとめたAwesomeリストです。

目次

• ライブラリ
• モデルとプロジェクト
• 動画
• 論文https://github.com/jax-ml/jax
• チュートリアルとブログ記事
• 書籍
• コミュニティ

• ニューラルネットワークライブラリ
  • Flax - 柔軟性と明確性を中心にした設計。
  • Flax NNX - 同じチームが開発したFlaxの進化形。
  • Haiku - シンプルさに焦点を当て、DeepMindのSonnetの開発者たちが作成した。 https://img.shields.io/github/stars/deepmind/dm-haiku?style=social”
  • Objax - PyTorchに類似したオブジェクト指向設計を採用。 https://img.shields.io/github/stars/google/objax?style=social”
  • Elegy - Deep Learning in JAX. Supports Flax, Haiku, and Optax. 向けのA High Level API。
  • Trax - common workloads. 向けの”Batteries included” deep learning library focused on providing solutions。
  • Jraph - 軽量なグラフニューラルネットワークライブラリ。 https://img.shields.io/github/stars/deepmind/jraph?style=social”
  • Neural Tangents - specifying neural networks of both finite and infinite width. 向けのHigh-level API。
  • HuggingFace Transformers - a wide range of natural language tasks (Flax). 向けのEcosystem of pretrained Transformers。
  • Equinox - 呼び出し可能なPyTreeとフィルタされたJIT/grad変換により、JAXにおけるニューラルネットワークを実現。 https://img.shields.io/github/stars/patrick-kidger/equinox?style=social”
  • Scenic - Computer Vision Research and Beyond. 向けのA Jax Library。
  • Penzai - 可読性・可視性・モデル編集の容易さを重視し、組み立て可能なツールとシンプルな認知モデルを用いる。 https://img.shields.io/github/stars/google-deepmind/penzai?style=social”
• Levanter - 名前付きテンソルとJAXを活用した、可読性・スケーラビリティ・再現性の高い基礎モデル。 https://img.shields.io/github/stars/stanford-crfm/levanter?style=social”
• EasyLM - LLMsを簡単に行う：JAX/FlaxでLLMの事前学習、微調整、評価、サービス。
• NumPyro - Pyroライブラリに基づく確率プログラミング。
• Chex - 信頼性の高いJAXコードの作成とテストに使うユーティリティ。
• Optax - 勾配処理と最適化ライブラリ。
• RLax - implementing reinforcement learning agents. 向けのLibrary。
• JAX, M.D. - 高速化された差分分子動力学。
• Coax - RL論文を簡単な方法でコードに変換。
• Distrax - TensorFlow Probabilityの再実装。確率分布と変換（bijectors）を含む。
• cvxpylayers - 微分可能な凸最適化層の構築。
• TensorLy - テンソルベースの学習をシンプルに。
• NetKet - Quantum Physics. 向けのMachine Learning toolbox。
• Fortuna - Uncertainty Quantification in Deep Learning. 向けのAWS library。
• BlackJAX - JAX. 向けのLibrary of samplers。
• Dynamax - 確率的状態空間モデル。

• ニューラルネットワークライブラリ
  • FedJAX - JAXにおけるフェデレートラーニング。OptaxとHaikuをベースに構築。
  • Equivariant MLP - 等変換性を持つニューラルネットワーク層の構築。
  • jax-resnet - ResNet variants in Flax. 向けのImplementations and checkpoints。
  • jax-raft - RAFT光学流推定器のJAX/Flaxへの移植。
  • Parallax - JAX. 向けのImmutable Torch Modules。
• 非線形最適化
  • Optimistix - 根の探索、最小化、固定点、最小二乗法。
  • JAXopt - GPU/TPUによるハードウェア加速、バッチ処理可能かつ微分可能な最適化アルゴリズム（JAXで）。
• jax-unirep - protein machine learning applications. <img src=” align=“center”>向けのLibrary implementing the 。 UniRep model https://img.shields.io/github/stars/ElArkk/jax-unirep?style=social”
• flowjax - 確率分布と正規化フローをEquinoxモジュールとして構築。
• flaxdiff - building and training Diffusion models in multi-node multi-device distributed settings (TPUs) 向けのFramework and Library。
• jax-flows - JAXにおける正規化フロー。
• sklearn-jax-kernels - scikit-learnカーネル行列のJAXでの利用。
• jax-cosmo - JAXにおける指数族分布。
• efax - JAXにおける可微分宇宙論ライブラリ。
• mpi4jax - CPUおよびGPUでのJAXコードにMPI操作を組み合わせる。
• imax - 画像の拡張および変換。
• FlaxVision - フラックス版TorchVision。
• Oryx - プログラム変換をベースとした確率プログラミング言語。
• Optimal Transport Tools - 最適輸送問題を解決するためのユーティリティをまとめるツールボックス。
• delta PV - 自動微分を用いた太陽電池シミュレータ。
• jaxlie - rigid body transformations and optimization. 向けのLie theory library。
• BRAX - these environments. 向けのDifferentiable physics engine to simulate environments along with learning algorithms to train agents。
• flaxmodels - Jax/Flax. 向けのPretrained models。
• CR.Sparse - sparse representations and compressive sensing. 向けのXLA accelerated algorithms。
• exojax - JAXと互換可能な系外行星・褐矮星の自動微分スペクトルモデル。
• PIX - JAX. 向けのPIX is an image processing library in JAX,。
• bayex - JAXを用いたベイズ最適化。
• JaxDF - differentiable simulators with arbitrary discretizations. 向けのFramework。
• tree-math - 配列に作用する関数を、PyTreeに作用する関数に変換。
• jax-models - 元々コードがなく、またはJAX以外のフレームワークで書かれた研究論文の実装。
• PGMax - building discrete Probabilistic Graphical Models (PGM’s) and running inference inference on them via JAX. 向けのA framework。
• EvoJAX - ハードウェア加速によるニューロ進化 https://img.shields.io/github/stars/google/evojax?style=social”
• evosax - JAXベースの進化戦略 https://img.shields.io/github/stars/RobertTLange/evosax?style=social”
• SymJAX - シンボリックなCPU/GPU/TPUプログラミング https://img.shields.io/github/stars/SymJAX/SymJAX?style=social”
• mcx - performant inference. 向けのExpress & compile probabilistic programs。
• Einshape - JAX and other frameworks. 向けのDSL-based reshaping library。
• ALX - distributed matrix factorization using Alternating Least Squares, more info in ALX: Large Scale Matrix Factorization on TPUs向けのOpen-source library。
• Diffrax - JAXにおける数値微分方程式ソルバー https://img.shields.io/github/stars/patrick-kidger/diffrax?style=social”
• tinygp - JAXにおける最も小さなガウス過程ライブラリ https://img.shields.io/github/stars/dfm/tinygp?style=social”
• gymnax - よく知られたgym APIを用いた強化学習環境 https://img.shields.io/github/stars/RobertTLange/gymnax?style=social”
• Mctx - JAXに内蔵されたモンテカルロ木探索アルゴリズム。
• KFAC-JAX - NNs. 向けのSecond Order Optimization with Approximate Curvature。
• TF2JAX - 関数／グラフをJAX関数に変換。
• jwave - differentiable acoustic simulations 向けのA library。
• GPJax - JAXにおけるガウス過程の実装。
• Jumanji - JAXで記述された、業界向けハードウェア加速型強化学習環境のソースコード。
• Eqxvision - TorchvisionのEquinox版。
• JAXFit - nonlinear least-squares problems (see arXiv paper). 向けのAccelerated curve fitting library。
• econpizza - JAXを用いて異質なエージェントを用いたマクロ経済モデルを解く。
• SPU - MPC（セキュアマルチパーティ計算）を用いたJAXコードを実行するためのドメイン特定コンパイラおよびランタイムセット。
• jax-tqdm - JAXのスキャンとループにtqdmの進行バーを追加する。
• safejax - JAX、Flax、Haiku、またはObjaxのモデルパラメータを🤗safetensorsでシリアル化する。
• Kernex - JAXにおける微分可能なステンルデコレータ。
• MaxText - Google CloudのTPUをターゲットとした、純粋なPython/JAXで書かれたシンプルで高性能かつスケーラブルなJaxLLM。
• Pax - training large scale models. 向けのA Jax-based machine learning framework。
• Praxis - Pax with a goal to be usable by other JAX-based ML projects. 向けのThe layer library。
• purejaxrl - JAXで実装されたベクトル化された、端到端の強化学習アルゴリズム。
• Lorax - JAXモデル（Flax、Haikuなど）にLoRAを自動的に適用。
• SCICO - JAXを用いた科学的計算画像処理。
• Spyx - machine learning on neuromorphic hardware. 向けのSpiking Neural Networks in JAX。
• 脳ダイナミクス・プログラミングエコシステム
  • BrainPy - Pythonで脳動力学プログラミング。
  • brainunit - JAXにおける物理単位と単位を意識した数学システム。
  • dendritex - JAXを用いた樹状細胞モデル化。
  • brainstate - Program Compilation and Augmentation. 向けのState-based Transformation System。
  • braintaichi - Taichi Langを活用した脳動力学演算子のカスタマイズ。
• OTT-JAX - JAXにおける最適輸送ツール。
• QDax - Jaxにおける品質多様性最適化。
• JAX Toolbox - JAX on NVIDIA GPUs using libraries such as T5x, Paxml, and Transformer Engine. 向けのNightly CI and optimized examples。
• Pgx - RL with an AlphaZero example. 向けのVectorized board game environments。
• EasyDeL - training and serving (Llama, MPT, Mixtral, Falcon, etc) in JAX 向けのEasyDeL 🔮 is an OpenSource Library to make your training faster and more Optimized With cool Options。
• XLB - Physics-Based Machine Learning. 向けのA Differentiable Massively Parallel Lattice Boltzmann Library in Python。
• dynamiqs - JAXを用いた量子系の高性能かつ微分可能なシミュレーション。
• foragax - JAXにおけるアグリーティブなモデリングフレームワーク。
• tmmax - thin-film optics research 向けのVectorized calculation of optical properties in thin-film structures using JAX. Swiss Army knife tool。
• Coreax - finding coresets to compress large datasets while retaining their statistical properties. 向けのAlgorithms。
• NAVIX - 強化学習環境MiniGridのJAXによる再実装。
• FDTDX - JAXを用いた有限差分時域電磁シミュレーション。
• DiffeRT - Radio Propagation powered by the JAX ecosystem. 向けのDifferentiable Ray Tracing toolbox。
• JAX-in-Cell - electron and ion dynamics in electromagnetic fields 向けのPlasma physics simulations using a PIC (Particle-in-Cell) method to self-consistently solve。
• kvax - JAX with support for efficient document mask computation and context parallelism. 向けのA FlashAttention implementation。
• astronomix - astrophysics in JAX 向けのdifferentiable (magneto)hydrodynamics。
• vivsim - 浸没境界-格子ボルツマン法を用いた流体-構造相互作用シミュレーション。
• MBIRJAX - 高性能トモグラフィ再構成。
• torchax - Jax to interoperate with model code written in PyTorch.向けのtorchax is a library。

• Fourier Feature Networks - の公式実装。<img src=” align=“center”> Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains
• kalman-jax - Markov (i.e., temporal) Gaussian processes using iterated Kalman filtering and smoothing向けのApproximate inference。
• jaxns - JAXにおけるネストサンプリング
• Amortized Bayesian Optimization - Amortized Bayesian Optimization over Discrete Spacesに関連するコード。
• Accurate Quantized Training - running and analyzing neural network quantization experiments in JAX and Flax向けのTools and libraries。
• BNN-HMC - What Are Bayesian Neural Network Posteriors Really Like?論文の実装。
• JAX-DFT - JAXにおける一次元密度関数理論（DFT）の実装、Kohn-Sham equations as regularizer: building prior knowledge into machine-learned physicsの実装。
• Robust Loss - A General and Adaptive Robust Loss Function論文の参照実装。
• Symbolic Functionals - Evolving symbolic density functionalsに関するデモンストレーション。
• TriMap - TriMap: Large-scale Dimensionality Reduction Using Tripletsの公式JAX実装。

Flax

• awesome-jax-flax-llms - JAX・Flaxで実装されたLLMのコレクション
• DeepSeek-R1-Flax-1.5B-Distill - DeepSeek-R1 1.5Bの蒸留推論LLMのFlax実装
• Performer - Performer（FAVOR+による線形トランスフォーマー）アーキテクチャのFlax実装
• JaxNeRF - マルチデバイスGPU／TPUをサポートするNeRF: Representing Scenes as Neural Radiance Fields for View Synthesisの実装。
• mip-NeRF - Mip-NeRF: A Multiscale Representation for Anti-Aliasing Neural Radiance Fieldsの公式実装。
• RegNeRF - RegNeRF: Regularizing Neural Radiance Fields for View Synthesis from Sparse Inputsの公式実装。
• JaxNeuS - NeuS: Learning Neural Implicit Surfaces by Volume Rendering for Multi-view Reconstructionの実装。
• Big Transfer (BiT) - Big Transfer (BiT): General Visual Representation Learningの実装。
• JAX RL - 強化学習アルゴリズムの実装
• gMLP - Pay Attention to MLPsの実装。
• MLP Mixer - MLP-Mixer: An all-MLP Architecture for Visionの最小実装。
• Distributed Shampoo - Second Order Optimization Made Practicalの実装。
• NesT - Aggregating Nested Transformersの公式実装
• XMC-GAN - Cross-Modal Contrastive Learning for Text-to-Image Generationの公式実装。
• FNet - FNet: Mixing Tokens with Fourier Transformsの公式実装
• GFSA - Learning Graph Structure With A Finite-State Automaton Layerの公式実装
• IPA-GNN - Learning to Execute Programs with Instruction Pointer Attention Graph Neural Networksの公式実装
• Flax Models - Flaxで実装されたモデルと手法のコレクション
• Protein LM - proteins, as described in Biological Structure and Function Emerge from Scaling Unsupervised Learning to 250 Million Protein Sequences and ProGen: Language Modeling for Protein Generation向けのImplements BERT and autoregressive models。
• Slot Attention - Differentiable Patch Selection for Image Recognition向けのReference implementation。
• Vision Transformer - An Image is Worth 16x16 Words: Transformers for Image Recognition at Scaleの公式実装。
• FID computation - mseitzer/pytorch-fidのFlaxへの移植
• ARDM - Autoregressive Diffusion Modelsの公式実装
• D3PM - Structured Denoising Diffusion Models in Discrete State-Spacesの公式実装
• Gumbel-max Causal Mechanisms - Learning Generalized Gumbel-max Causal Mechanisms, with extra code in GuyLor/gumbel_max_causal_gadgets_part2向けのCode。
• Latent Programmer - ICML 2021のLatent Programmer: Discrete Latent Codes for Program Synthesis論文のコード。
• SNeRG - Baking Neural Radiance Fields for Real-Time View Synthesisの公式実装。
• Spin-weighted Spherical CNNs - Spin-Weighted Spherical CNNsの適応
• VDVAE - Very Deep VAEs Generalize Autoregressive Models and Can Outperform Them on Imagesの適応。オリジナルコードはopenai/vdvaeに存在
• MUSIQ - ICCV 2021のMUSIQ: Multi-scale Image Quality Transformer論文向けチェックポイントとモデル推論コード。
• AQuaDem - Continuous Control with Action Quantization from Demonstrationsの公式実装
• Combiner - Combiner: Full Attention Transformer with Sparse Computation Costの公式実装
• Dreamfields - the ICLR 2022 paper Progressive Distillation for Fast Sampling of Diffusion Modelsの公式実装。
• GIFT - Gradual Domain Adaptation in the Wild:When Intermediate Distributions are Absentの公式実装
• Light Field Neural Rendering - Light Field Neural Renderingの公式実装
• Sharpened Cosine Similarity in JAX by Raphael Pisoni - JAX/Flax による鋭いコサイン類似度層の実装
• GNNs for Solving Combinatorial Optimization Problems - JAX + FlaxによるCombinatorial Optimization with Physics-Inspired Graph Neural Networksの実装
• DETR - Sinkhornソルバーと並列双対マッチングを用いたDETR: End-to-end Object Detection with TransformersのFlax実装

Haiku

• AlphaFold - AlphaFold v2.0で提示されたHighly accurate protein structure prediction with AlphaFoldの推論パイプライン実装
• Adversarial Robustness - Uncovering the Limits of Adversarial Training against Norm-Bounded Adversarial Examples and Fixing Data Augmentation to Improve Adversarial Robustness向けのReference code。
• Bootstrap Your Own Latent - Bootstrap your own latent: A new approach to self-supervised Learning論文の実装。
• Gated Linear Networks - GLNsは、バックプロパゲーションなしのニューラルネットワークのファミリーである
• Glassy Dynamics - Unveiling the predictive power of static structure in glassy systems論文のオープンソース実装。
• MMV - the models in Self-Supervised MultiModal Versatile Networks向けのCode。
• Normalizer-Free Networks - NFNetsの公式Haiku実装
• NuX - JAXにおけるノーマライズフローの実装
• OGB-LSC - DeepMindがPCQM4M-LSC（量子化学）とMAG240M-LSC（学術グラフ）を OGB Large-Scale Challenge（OGB-LSC）へ提出した実装。
• Persistent Evolution Strategies - Unbiased Gradient Estimation in Unrolled Computation Graphs with Persistent Evolution Strategies論文で使用したコード。
• Two Player Auction Learning - Auction learning as a two-player game論文のJAX実装。
• WikiGraphs - WikiGraphs: A Wikipedia Text - Knowledge Graph Paired Dataseの結果を再現するベースラインコード。

Trax

• Reformer - Reformer（効率的なトランスフォーマー）アーキテクチャの実装

NumPyro

• lqg - linear-quadratic Gaussian problems from the paper Putting perception into action with inverse optimal control for continuous psychophysics向けのOfficial implementation of Bayesian inverse optimal control。

Equinox

• Sampling Path Candidates with Machine Learning - Towards Generative Ray Path Sampling for Faster Point-to-Point Ray Tracing論文の公式チュートリアル兼実装。

• NeurIPS 2020: JAX Ecosystem Meetup - JAX、DeepMindでの利用、およびエンジニア、科学者、JAXコアチームとのディスカッション
• Introduction to JAX - JAXでゼロからシンプルなニューラルネットワークの実装
• JAX: Accelerated Machine Learning Research | SciPy 2020 | VanderPlas - JAXの核心設計、新たな研究への貢献、およびその利用を開始する方法について。
• Bayesian Programming with JAX + NumPyro — Andy Kitchen - NumPyroを用いたベイズモデルの導入講座。
• JAX: Accelerated machine-learning research via composable function transformations in Python | NeurIPS 2019 | Skye Wanderman-Milne - Program Transformations for Machine LearningワークショップでのJAX入門発表。
• JAX on Cloud TPUs | NeurIPS 2020 | Skye Wanderman-Milne and James Bradbury - TPUホストアクセスの紹介とデモ。
• Deep Implicit Layers - Neural ODEs, Deep Equilibirum Models, and Beyond | NeurIPS 2020 - Zico Kolter、David Duvenaud、Matt Johnson が作成したチュートリアル．Colabノートブックは Deep Implicit Layers に公開されている．
• Solving y=mx+b with Jax on a TPU Pod slice - Mat Kelcey - v3-32 TPU Podスライス上でデータ並列アプローチによる訓練まで進む、YouTubeチュートリアルシリーズ（Colabノートブック付き）。
• JAX, Flax & Transformers 🤗 - JAX／Flax、Transformer、大規模言語モデルなど、素晴らしいテーマを中心に3日間の講演会。

この節にはJAXを主題とする論文（JAXベースのライブラリのホワイトペーパーやJAX自体の研究など）を掲載します。JAXで実装された論文はモデル／プロジェクト節に掲載します。

• Compiling machine learning programs via high-level tracing. Roy Frostig, Matthew James Johnson, Chris Leary. MLSys 2018. - JAXの初期バージョンについての白書。計算のトレースとコンパイル方法を詳細に説明。
• JAX, M.D.: A Framework for Differentiable Physics. Samuel S. Schoenholz, Ekin D. Cubuk. NeurIPS 2020. - JAX、M.D.を導入する。これはシミュレーション環境、相互作用ポテンシャル、ニューラルネットワークなどを含む微分可能物理ライブラリである。
• Enabling Fast Differentially Private SGD via Just-in-Time Compilation and Vectorization. Pranav Subramani, Nicholas Vadivelu, Gautam Kamath. arXiv 2020. - JAXのJITとVMAPを活用し、既存のライブラリよりも高速な差分プライバシーを達成する。
• XLB: A Differentiable Massively Parallel Lattice Boltzmann Library in Python. Mohammadmehdi Ataei, Hesam Salehipour. arXiv 2023. - XLBライブラリに関する白書：ベンチマーク、検証、ライブラリの詳細情報について。

• Using JAX to accelerate our research by David Budden and Matteo Hessel - DeepMindにおけるJAXおよびJAXエコシステムの現状について説明する。
• Getting started with JAX (MLPs, CNNs & RNNs) by Robert Lange - 基本的なJAX演算子から始めて、ニューラルネットワークの構成ブロックを構築する。
• Learn JAX: From Linear Regression to Neural Networks by Rito Ghosh - JAXの基礎を丁寧に解説し、線形回帰やロジスティック回帰、ニューラルネットワークモデルを実装し、現実世界の問題に応用する方法を紹介する。
• Tutorial: image classification with JAX and Flax Linen by 8bitmp3 - FlaxのLinen APIを使ってシンプルなコネクショナルネットワークを作成し、手書き数字を認識できるように訓練する方法を学習する。
• Plugging Into JAX by Nick Doiron - Kaggleの花の分類チャレンジにおいて、Flax、Haiku、Objaxの比較を紹介する。
• Meta-Learning in 50 Lines of JAX by Eric Jang - JAXとメタ学習についての概要を解説する。
• Normalizing Flows in 100 Lines of JAX by Eric Jang - RealNVP の簡潔な実装．
• Differentiable Path Tracing on the GPU/TPU by Eric Jang - パスティングの実装に関するチュートリアル。
• Ensemble networks by Mat Kelcey - アンサンブルネットは、複数のモデルを1つの論理モデルとして表現する手法である。
• Out of distribution (OOD) detection by Mat Kelcey - OOD detection向けのImplements different methods。
• Understanding Autodiff with JAX by Srihari Radhakrishna - JAXを使って自動微分がどのように機能するかを理解する。
• From PyTorch to JAX: towards neural net frameworks that purify stateful code by Sabrina J. Mielke - パイトorch風のコードスタイルから機能型コードスタイルへの移行方法を示す。
• Extending JAX with custom C++ and CUDA code by Dan Foreman-Mackey - JAXでカスタムオペレーションを提供するためのインフラストラクチャを解説するチュートリアル。
• Evolving Neural Networks in JAX by Robert Tjarko Lange - JAXがスケーラブルなニューローエволユーションアルゴリズムの次世代を支える方法を検証する。
• Exploring hyperparameter meta-loss landscapes with JAX by Luke Metz - JAXを用いて、内側損失の最適化（SGDおよびモーメンタム）、外側損失の勾配最適化、および進化戦略による外側損失最適化を行う方法を示す。
• Deterministic ADVI in JAX by Martin Ingram - JAXを用いて、自動微分変分推論（ADVI）を簡単にかつクリーンに実装するプロセスを解説する。
• Evolved channel selection by Mat Kelcey - a particular loss向けのTrains a classification model robust to different combinations of input channels at different resolutions, then uses a genetic algorithm to decide the best combination。
• Introduction to JAX by Kevin Murphy - 言語のさまざまな側面を紹介し、簡単な機械学習問題に応用するColab。
• Writing an MCMC sampler in JAX by Jeremie Coullon - JAXでMCMCサンプラーを書くさまざまな方法についてのチュートリアルと、性能ベンチマーク。
• How to add a progress bar to JAX scans and loops by Jeremie Coullon - JAXにおけるコンパイルされたループに進行バーを追加する方法についてのチュートリアル． host_callback モジュールを使用する．
• Get started with JAX by Aleksa Gordić - ゼロからJAXの知識を身につけ、Haikuでニューラルネットワークを構築するまでのノートブックと動画シリーズ。
• Writing a Training Loop in JAX + FLAX by Saurav Maheshkar and Soumik Rakshit - JAX、Flax、Optaxを使ってシンプルな端到端のトレーニングと評価パイプラインを書くチュートリアル。
• Implementing NeRF in JAX by Soumik Rakshit and Saurav Maheshkar - JAXでニューラルラジエンスフィールドによって表現されたシーンの3D体積レンダリングについてのチュートリアル。
• Deep Learning tutorials with JAX+Flax by Phillip Lippe - 基礎から最新のディープラーニング概念（例：JAX/Flaxの紹介、活性化関数）まで、PyTorchへの翻訳付きのさまざまな深層学習概念を解説するノートブックシリーズ。
• Achieving 4000x Speedups with PureJaxRL - JAXがベクトル化によって強化学習のトレーニングを大幅に加速する方法についてのブログ記事。
• Simple PDE solver + Constrained Optimization with JAX by Philip Mocz - JAXを用いて拡散伝導方程式を解く簡単な例と、その結果を制約付き最適化問題に適用して、望ましい結果をもたらす初期条件を求める方法。

• Jax in Action - JAXを体系的に学ぶための書籍です。

• JaxLLM (Unofficial) Discord
• JAX GitHub Discussions
• Reddit

コントリビューション

コントリビューションを歓迎します。最初にコントリビューションガイドラインをお読みください。