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.

An Extensible Deep Learning Library

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