fermix

Fast batched determinants and Pfaffians in JAX

Fermix supports slogdet / slogpf / det / pf for float32, float64, complex64, and complex128 matrices on NVIDIA GPUs, written in JAX with Pallas (Triton) kernels. The codes are generated by Claude Code.

Designed for quantum Monte Carlo in fermionic systems: a large batch of moderate-size matrices (n ~ 32-8192), forward and backward.

Install

Requires jax >= 0.7.1

pip install fermix

Usage

import jax, jax.numpy as jnp
from fermix import slogdet, slogpf, det, pf

A = jax.random.normal(jax.random.key(0), (4096, 128, 128), jnp.float32)
sign, logabs = slogdet(A) # like jnp.linalg.slogdet, but faster
S = A - jnp.swapaxes(A, -1, -2)
sign, logabs = slogpf(S) # Pfaffian of the skew-symmetric batch

Simply replace slogdet / slogpf / det / pf in your program by the functions in fermix.

All functions accept any leading batch dimensions, any matrix size, the four dtypes float32 / float64 / complex64 / complex128 (the 64-bit ones need jax_enable_x64), and are jit/vmap/grad/jvp compatible. Conventions follow jnp.linalg.slogdet.

Performance

Forward time per matrix on one H200 GPU against jnp.linalg.slogdet and a naive batched Parlett–Reid slogpf written in jax.numpy (the generic fallback path). The time costs of det and pf are similar.

Notes and limits

• Without a CUDA GPU the call emits a FermixFallbackWarning and runs a much slower generic jax.numpy path; dtypes other than the four supported ones raise TypeError.
• Launch parameters are tuned per GPU architecture and dtype (fermix._common.TUNES); only the A100-80GB and the H200 have measured tables, other GPUs use the closest one.
• prec selects the block-GEMM algorithm: "tf32x3" (default, tensor cores at fp32-level accuracy) or "ieee" (exact fp32, ~5 % slower); the 64-bit dtypes always run IEEE fp64.
• Complex matrices are stored as separate real and imaginary arrays (Triton has no complex type), so they move twice the bytes and do four real products per complex one.
• Memory: two (B, n, n) work buffers for the forward and about five more for the gradient.
• Compile time is a few seconds per shape and dtype up to n ~ 1024 and tens of seconds to minutes for n in the thousands; JAX's persistent compilation cache (jax_compilation_cache_dir) makes it a one-off.
• det and pf overflow the 32-bit dtypes for large or badly scaled matrices, and so do their gradients; use slogdet and slogpf there.
• Singular inputs are safe: the forward gives (0, −inf) and the gradients stay finite, never NaN.

Citation

@software{fermix,
  author  = {Chen, Ao},
  title   = {{fermix}: Fast batched determinants and Pfaffians in {JAX}},
  year    = {2026},
  version = {0.1.0},
  url     = {https://github.com/ChenAo-Phys/fermix}
}