Xla Python Packages

torch-xla

Enabling PyTorch on XLA Devices (e.g. Google TPU)

dense-evolution

NISQ Quantum Simulation toolkit- JAX-Native.

mpi4jax

Zero-copy MPI communication of JAX arrays, for turbo-charged HPC applications in Python :zap:

jax-tpu-embedding

High-performance large-scale embedding acceleration for JAX on Google TPU SparseCores.

jaxdecomp

JAX reimplementation and bindings for the NVIDIA cuDecomp library

strix-halo-jax

Run JAX and MuJoCo MJX on AMD Ryzen AI MAX (Strix Halo, gfx1151) from pip wheels - no system ROCm - with a hang-safe check that tells you which wheel set works

revlib

Simple and efficient RevNet-Library for PyTorch with XLA and DeepSpeed support and parameter offload

felafax

Felafax is building AI infra for non-NVIDIA GPUs

maddening

Modular Automatic Differentiation and Data-Enhanced Neural-network INteracting Graph (MADDENING). A pure JAX end-to-end autodifferentiable graph based framework for multiphysics simulation.

gpu-clahe

GPU-accelerated CLAHE for TensorFlow: ~15x faster than single-threaded OpenCV, with bit-identical output across CPU and GPU

tmmax

A fast transfer matrix method written in jax for modelling optical multilayer thin films

jaxplorer

A compiler explorer TUI for JAX.

pypjrt

A dependency-free Python client for the PJRT C API and the XLA FFI. Run compiled ML programs on CPUs, GPUs and TPUs from pure Python.

tf-q-finance

High-performance TensorFlow library for quantitative finance.

tl-pjrt

Tensor Logic as a StableHLO-emitting DSL for PJRT runtimes

katmer

katmer is a powerful library for optimizing the design of optical thin films using automatic differentiation via JAX and Equinox, enabling efficient and accurate inverse design solutions.