Graph Attention Networks (GATs) for node classification and regression tasks

Searching for galaxy satellites by their impact on lens potential in strong gravitational lensing setup

Kalman filter in JAX

Comparative implementation of Classical and Quantum Physics-Informed Neural Networks (PINNs) for the 2D Navier-Stokes cylinder wake problem, using JAX and PennyLane.

Translate SQL with AI

Implementation of FastSpeech2 in JAX

API for recipe generation from flax-community/t5-recipe-generation (huggingface)

This repository provides a combination of the bubble-point algorithm and Naphtali-Sandholm algorithm to steadily compute distillation separations with partial condensers

a flexible reinforcement learning playground for turn-based X×X board games. Two RL agents train by playing against each other (self-play), learning strategies to win purely from rewards and game outcomes.

World Models (Ha & Schmidhuber) reproduced in JAX/Equinox. Features parallelized 'dream' training (256 agents), asymmetric loss for Sim2Real transfer, and active failure learning.

One-command determinism + manifest for ML projects.

Google Summer of Code 2024

A JAX / NNX implementation of a VQ-VAE for audio compression

Unified Machine Learning/Deep Learning demo repository combining most of the popular ML frameworks.

A generalized implementation of Grad-CAM for Flax

fcpe jax flax version

Add ppc64le support back to BoringSSL for use with XLA

This tutorial is part of the "Intensive Course Machine Learning: From Basics to Advanced Concepts" of the QuCoLiMa research network and the International Max Planck Research School Physics of Light.

Accelerated Quantum Reinforcement Learning Benchmarking with JAX