jrmt — JAXed Random Matrix Theory

A JAX port of thrmt. Pure functional samplers for the classical real / complex / quantum random matrix ensembles, built on jax.numpy and jaxtyping.

Parity with thrmt

jrmt mirrors thrmt's public surface 1:1 — all 25 names from thrmt's __init__ (16 ensembles + 9 historical aliases) are available under jrmt. The single intentional API break is described below.

For random_rho_pure, thrmt's bo_einsum keyword is gone — the JAX implementation always uses jnp.einsum for the outer product.

Install

pip install -e .

Requires Python ≥ 3.10, jax, and jaxtyping. Enable 64-bit precision before importing anything that uses complex128:

from jax import config
config.update("jax_enable_x64", True)

Use

import jax
import jax.numpy as jnp
import jrmt

key = jax.random.key(0)

# Single draw
h = jrmt.random_gue(key, size=8, sigma=1.0)
assert jnp.allclose(h, jnp.conjugate(h.T))

# Batched draw via vmap
keys = jax.random.split(key, 1000)
batch = jax.vmap(lambda k: jrmt.random_rho_hs(k, size=8))(keys)
assert batch.shape == (1000, 8, 8)

Tests

pytest                 # fast invariants + cross-library moment-matching
pytest -m slow         # asymptotic-law KS tests (~30 s)

Cross-library moment-matching (tests/test_cross_lib.py) needs both thrmt and torch installed in the environment; the other suites do not.