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.
| Migration of … | thrmt (PyTorch) |
jrmt (JAX) |
|---|---|---|
| PRNG | implicit, global torch.manual_seed(...) |
explicit, key is the first positional argument of every random function |
| Batching | batch_shape=(B,) kwarg returning (B, …) |
jax.vmap(lambda k: fn(k, …))(jax.random.split(root_key, B)) — pure-functional |
| Device | device=... kwarg |
dropped — use jax.device_put or JAX's implicit placement |
| dtype | torch.{cdouble, cfloat, double, float} |
jnp.{complex128, complex64, float64, float32, bfloat16, float16} |
| Shape types | runtime checks at the wrapper | jaxtyping annotations: Float[Array, "n n"], Complex[Array, "n n"], PRNGKeyArray |
Output of random_obs_csu's evdist callback |
evdist(*batch_shape, size, dtype, device) |
evdist(key, size, dtype) — the only documented public-API break (forced by JAX's PRNG plumbing) |
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.
1 maintainer
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
File details
Details for the file jrmt-0.1.0.tar.gz.
File metadata
- Download URL: jrmt-0.1.0.tar.gz
- Upload date:
- Size: 17.3 kB
- Tags: Source
- Uploaded using Trusted Publishing? Yes
- Uploaded via:
twine/6.1.0 CPython/3.13.12
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
0ce5edcb9ec1cf689e7cc42db7da151b24d5669a0b7e6c90e2add4810ed90f7e
|
|
| MD5 |
7b9d81015718d2951e0c2304ad410278
|
|
| BLAKE2b-256 |
d34f8195b808f7e3b2ec44c52349f6e513f4eabb4eb11b6e2deb231e2b8fadac
|
Provenance
The following attestation bundles were made for jrmt-0.1.0.tar.gz:
Publisher:
pypipublish.yml on emaballarin/jrmt
-
Statement:
-
Statement type:
https://in-toto.io/Statement/v1 -
Predicate type:
https://docs.pypi.org/attestations/publish/v1 -
Subject name:
jrmt-0.1.0.tar.gz -
Subject digest:
0ce5edcb9ec1cf689e7cc42db7da151b24d5669a0b7e6c90e2add4810ed90f7e - Sigstore transparency entry: 1521582332
- Sigstore integration time:
-
Permalink:
emaballarin/jrmt@31fe0176549d826ab1f090af6c928c42f0fcfac4 -
Branch / Tag:
refs/tags/v0.1.0 - Owner: https://github.com/emaballarin
-
Access:
private
-
Token Issuer:
https://token.actions.githubusercontent.com -
Runner Environment:
github-hosted -
Publication workflow:
pypipublish.yml@31fe0176549d826ab1f090af6c928c42f0fcfac4 -
Trigger Event:
push
-
Statement type:
File details
Details for the file jrmt-0.1.0-py3-none-any.whl.
File metadata
- Download URL: jrmt-0.1.0-py3-none-any.whl
- Upload date:
- Size: 12.5 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? Yes
- Uploaded via:
twine/6.1.0 CPython/3.13.12
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
68b59f3fae20cda9e58956870d04f9cd932ab47c84f5e2f6fd8953298e5df91d
|
|
| MD5 |
3ae2f39309f47543942400d81e257530
|
|
| BLAKE2b-256 |
c5d17656a2c06011056e778f622f615a1738696e583a1619386fe751689ea239
|
Provenance
The following attestation bundles were made for jrmt-0.1.0-py3-none-any.whl:
Publisher:
pypipublish.yml on emaballarin/jrmt
-
Statement:
-
Statement type:
https://in-toto.io/Statement/v1 -
Predicate type:
https://docs.pypi.org/attestations/publish/v1 -
Subject name:
jrmt-0.1.0-py3-none-any.whl -
Subject digest:
68b59f3fae20cda9e58956870d04f9cd932ab47c84f5e2f6fd8953298e5df91d - Sigstore transparency entry: 1521582343
- Sigstore integration time:
-
Permalink:
emaballarin/jrmt@31fe0176549d826ab1f090af6c928c42f0fcfac4 -
Branch / Tag:
refs/tags/v0.1.0 - Owner: https://github.com/emaballarin
-
Access:
private
-
Token Issuer:
https://token.actions.githubusercontent.com -
Runner Environment:
github-hosted -
Publication workflow:
pypipublish.yml@31fe0176549d826ab1f090af6c928c42f0fcfac4 -
Trigger Event:
push
-
Statement type:
