Skip to main content
Context Awesome

awesome-jax

github.com/lockwo/awesome-jax ↗

Curated list of JAX Resources and Packages

43
GitHub Stars
149
Curated Resources
4
Categories
2 hours ago
Last Refreshed
LibrariesModels and ProjectsTutorials and Blog PostsCommunity

Use this list with your AI agent

Add the Context Awesome MCP server to Claude, Cursor, or any MCP client, then ask:

"Show me videos resources from awesome-jax"

Installation instructions →

What's inside

Libraries

  • astronomix

    differentiable (magneto)hydrodynamics for astrophysics in JAX.

  • bayex

    Minimal Implementation of Bayesian Optimization in JAX.

  • blackjax

    BlackJAX is a Bayesian Inference library designed for ease of use, speed and modularity.

  • brax

    Massively parallel rigidbody physics simulation on accelerator hardware.

  • cd_dynamaxUp and Coming Libraries

    Extension of dynamax repo to cases with continuous-time dynamics with measurements sampled at possibly irregular discrete times. Allows generic inference of dynamical systems parameters from partial noisy observations via auto-differentiable filtering, SGD, and HMC.

  • chex

    a library of utilities for helping to write reliable JAX code.

Models and Projects

  • esm2quinox

    An implementation of ESM2 in Equinox+JAX.

  • joltz

    JAX translation of boltz.

  • mosaic

    composite-objective protein design.

  • nequip-eqx

    JAX implementation of the NequIP neural network interatomic potential.

  • whisper-jax

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

Showing a sample of 149 resources. View the full list on GitHub →