VizTracer is a low-overhead logging/debugging/profiling tool that can trace and visualize your python code execution.

JAX accelerated Quantum Monte Carlo

Conic solver for optimization problems in quantum information theory

Matrix-free linear algebra in JAX.

JAX Synergistic Memory Inspector

Turn jitted jax functions back into python source code

Jax-based quantum state vector simulator.

(Crafter + NetHack) in JAX. ICML 2024 Spotlight.

Tensor Parallelism with JAX + Shard Map

Bind any function written in another language to JAX with support for JVP/VJP/batching/jit compilation

This repository is maintained to release dataset and models for multimodal puzzle reasoning.

JAX for Graphcore IPU (experimental)

TessellateIPU: low level Poplar tile programming from Python

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

Memray is a memory profiler for Python

CUDA Templates for Linear Algebra Subroutines

Implementation of Forward Laplacian algorithm in JAX

Forward mode laplacian implemented in JAX tracer

A library for simulating a quantum circuit. Built initially to be a tutorial resource

Run PyTorch in JAX. 🤝

XLB: Accelerated Lattice Boltzmann (XLB) for Physics-based ML

Second Order Optimization and Curvature Estimation with K-FAC in JAX.