Metadata-Version: 2.4
Name: fermix
Version: 0.1.0
Summary: Fast batched real and complex determinants and Pfaffians on GPU with JAX/Pallas kernels, with singular-safe gradients
Author-email: Ao Chen <aochen@caltech.edu>
License-Expression: MIT
Project-URL: Homepage, https://github.com/ChenAo-Phys/fermix
Project-URL: Repository, https://github.com/ChenAo-Phys/fermix
Project-URL: Issues, https://github.com/ChenAo-Phys/fermix/issues
Keywords: jax,pallas,determinant,pfaffian,gpu,variational monte carlo
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering :: Physics
Classifier: Environment :: GPU :: NVIDIA CUDA
Classifier: Operating System :: POSIX :: Linux
Requires-Python: >=3.11
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: jax>=0.7.1
Requires-Dist: numpy
Provides-Extra: test
Requires-Dist: pytest; extra == "test"
Provides-Extra: dev
Requires-Dist: pytest; extra == "dev"
Requires-Dist: black; extra == "dev"
Requires-Dist: pyright; extra == "dev"
Dynamic: license-file

<h1 align='center'>fermix</h1>

<p align="center">
  <strong>Fast batched determinants and Pfaffians in JAX</strong>
</p>

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

```bash
pip install fermix
```

## Usage

```python
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.

![slogdet forward time and speedup over jnp.linalg.slogdet, H200](https://raw.githubusercontent.com/ChenAo-Phys/fermix/main/plots/benchmark_H200_slogdet.png)

![slogpf forward time and speedup over the naive jax.numpy path, H200](https://raw.githubusercontent.com/ChenAo-Phys/fermix/main/plots/benchmark_H200_slogpf.png)

## 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

```bibtex
@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}
}
```
