chore(release): cut v1.14.0 — GPU-acceleration sprint

Bumps version 1.13.1 → 1.14.0.

Adds the v1.14.0 CHANGELOG entry covering the 5-commit GPU sprint
landed earlier this session:
- 54d83b4  feat(neural): opt-in GPU/MPS routing via STATSPAI_TORCH_DEVICE
- 2fc18e4  feat(fast): Rust cluster_meat kernel with Rayon over clusters
- c339063  feat(iv):   sp.iv(absorb=...) HDFE Phase 3 — 2SLS with absorbed FEs
- 1ac1ba3  feat(fast): JAX-backed feols_jax for end-to-end GPU/TPU OLS
- b7bae77  feat(fast): vmap'd JAX bootstrap (pairs + cluster)
- 4782152  feat(fast): Phase 4c — wild + wild_cluster bootstrap (score form)
- 39d6ced  docs(gpu):  GPU acceleration guide + README pointer

Also extends paper.md (JSS) with a fifth bullet under *Unique features*
documenting the accelerator story; cites cameron2008bootstrap
(verified DOI in paper.bib) for the wild cluster bootstrap and notes
the new Rust HDFE / cluster-meat kernel in the implementation
paragraph. The CHANGELOG headline + Verified section together make
the v1.14 GPU promise explicit and bounded — most StatsPAI
estimators remain CPU-only by design.

Verified
- 60/60 v1.14 test bundle pass
  (test_torch_device_resolver, test_iv_absorb, test_jax_feols,
   test_jax_feols_bootstrap; test_cluster_meat_rust skips without
   maturin built).
- ``python -c "import statspai as sp; print(sp.__version__)"``
  reports 1.14.0; ``sp.fast.feols_jax_bootstrap`` resolves cleanly.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: brycewang-stanford

CHANGELOG.md

@@ -2,6 +2,117 @@
 
 All notable changes to StatsPAI will be documented in this file.
 
+## [1.14.0] — 2026-05-05
+
+### Headline
+
+GPU-acceleration sprint. Three workloads now opt into accelerator
+backends without changing their public API: (1) the neural causal
+estimators (`sp.deepiv`, `sp.tarnet`, `sp.cfrnet`, `sp.dragonnet`,
+`sp.cevae`) route through PyTorch CUDA / MPS via a centralised
+device resolver; (2) `sp.fast.feols_jax` runs the full WLS solve on
+JAX / XLA; and (3) `sp.fast.feols_jax_bootstrap` lifts a JIT-compiled
+single-iteration WLS kernel to a `jax.vmap` batched primitive,
+giving a 10–100x speedup over sequential CPU bootstrap on CUDA / TPU
+at B ≥ 1000. Four bootstrap variants share the same JAX kernel
+infrastructure: pairs (multinomial-weight resampling), cluster
+(Cameron–Gelbach–Miller 2008 §III.A), wild (row-level Rademacher),
+and wild cluster (Cameron–Gelbach–Miller 2008 §III.B); the wild
+variants use the score formulation
+`β* = β̂ + (X'WX)⁻¹ X'W (η ⊙ û)` which is mathematically identical
+to refitting on `y* = X β̂ + η ⊙ û` but needs one mat-vec per
+iteration instead of a full QR. A new `cluster_meat` Rust kernel
+in `statspai_hdfe` (PyO3 + Rayon, parallel over clusters) is wired
+behind `statspai.core._numba_kernels.cluster_meat` with the existing
+numba kernel as automatic fallback. `sp.iv(absorb=...)` is the new
+2SLS-with-HDFE entry point: residualises `y`, exogenous controls,
+endogenous regressors, and instruments by one or more FE columns
+via the Phase-1 Rust demean kernel before fitting, with the residual
+DOF adjusted by `Σ(G_k - 1)` to charge the absorbed FE rank against
+iid / HC1 / CR1 SEs. A new `docs/guides/gpu_acceleration.md` is the
+canonical landing page for the accelerator story; the README and
+`paper.md` link to it and explicitly bound the GPU promise (most
+estimators are CPU-only by design — DiD / RD / synth / GMM are
+bandwidth-bound or small-K convex programs where a tuned CPU
+kernel matches GPU performance).
+
+### Added
+
+- `sp.fast.feols_jax` — JAX-backed end-to-end OLS / WLS with HDFE.
+  Same formula DSL and `FeolsResult` return type as `sp.fast.feols`;
+  the WLS solve and HC1 sandwich run on the default JAX device.
+  CR1 cluster sandwich delegates to the existing `crve` (which
+  itself dispatches to the new `cluster_meat` Rust kernel when
+  built). Default `dtype="float64"` preserves bit-comparable
+  numerics; `dtype="float32"` available for the GPU fast path.
+- `sp.fast.feols_jax_bootstrap` — vmap'd bootstrap with four
+  variants (`pairs`, `cluster`, `wild`, `wild_cluster`).
+  `vmap_chunk_size` parameter for memory control on tight devices.
+  Same-seed → bit-identical reproducibility via `jax.random` PRNG.
+  Returns a `FeolsBootstrapResult` dataclass with `coef`, `se_boot`,
+  percentile `ci_lower` / `ci_upper`, and the full `boot_betas`
+  table for custom CI methods.
+- `sp.iv(absorb=...)` — 2SLS with HDFE residualisation. Accepts
+  `"firm + year"` string syntax or `["firm", "year"]` list.
+  LIML / Fuller / GMM / JIVE raise `NotImplementedError` (Phase 3b).
+- `STATSPAI_TORCH_DEVICE` environment variable (`cpu` / `cuda` /
+  `cuda:N` / `mps` / `auto`) routes neural causal estimators
+  through the requested device. Default `cpu` preserves existing
+  pinned numerics; explicit `cuda` raises if the device is
+  unavailable rather than silently falling back. New
+  `sp.fast.torch_device_info()` mirrors `sp.fast.jax_device_info()`.
+- `statspai_hdfe::cluster_meat` Rust kernel — Rayon parallel over
+  clusters with thread-local k×k upper-triangle accumulator and
+  elementwise reduction. Bumped the crate version 0.5.0-alpha.1 →
+  0.7.0-alpha.1. Activation requires a one-time
+  `pip install maturin && cd rust/statspai_hdfe && maturin develop --release`;
+  Python falls back to the numba kernel transparently when the
+  Rust extension is absent.
+- `docs/guides/gpu_acceleration.md` — accelerator landing page
+  with activation recipes, a Google Colab quickstart benchmark,
+  and an explicit "what is *not* GPU-accelerated and why" table.
+
+### Changed
+
+- `paper.md` adds a fifth bullet to the *Unique features*
+  list documenting the accelerator story, and notes the Rust
+  HDFE / cluster-meat kernel in the implementation paragraph.
+- `README.md` comparison-table accelerator row now links to the
+  new GPU guide; the *What StatsPAI is — and is not* bullet
+  expands to explicitly mention `feols_jax`,
+  `feols_jax_bootstrap`, and the vmap mechanism.
+
+### Internal
+
+- New helper `_jax_prep_inputs` shares formula-parse + FE-residualise
+  logic between `feols_jax` and `feols_jax_bootstrap`.
+  `feols_jax` itself is unchanged in this release; consolidation
+  into a shared call site is a candidate follow-up.
+- Rust crate adds `src/cluster.rs` (kernel) and a `cluster_meat`
+  PyO3 binding in `src/lib.rs`. 3 cargo unit tests cover small-DGP
+  reference parity, k=1 closed form, and empty-input safety.
+
+### Verified
+
+- 10 PyTorch device-resolver tests
+  (`tests/test_torch_device_resolver.py`); 51 existing neural
+  tests pass without numerical drift on default CPU.
+- 9 `cluster_meat` Rust parity tests
+  (`tests/test_cluster_meat_rust.py`) — auto-skip when
+  `statspai_hdfe` is not built.
+- 13 `sp.iv(absorb=)` parity tests vs explicit drop-first
+  dummies (`tests/test_iv_absorb.py`); coefficients agree to
+  `atol=1e-9`, iid SE to `rtol=1e-3`, cluster SE to `rtol=1e-2`.
+- 12 `feols_jax` parity tests vs `feols`
+  (`tests/test_jax_feols.py`); iid / hc1 / cr1 / weighted /
+  float32 / 6 error-path validations.
+- 24 `feols_jax_bootstrap` tests
+  (`tests/test_jax_feols_bootstrap.py`); convergence to HC1 SE
+  for pairs / wild and CR1 SE for cluster / wild_cluster at
+  B=2000 (rtol 10–15%); algebraic identity check that the wild
+  score formulation reproduces the literal "refit on pseudo-y"
+  bootstrap bit-for-bit on a no-FE DGP (`atol=1e-9`).
+
 ## [1.13.1] — 2026-05-05
 
 ### Headline

paper.md

@@ -234,11 +234,38 @@ interface: `.summary()`, `.plot()`, `.to_latex()`, `.to_docx()`, and
    progress notifications, and `sp.citation()` together with a
    `paper.bib`-checked DOI verifier that refuses to emit citations
    not present in the curated bibliography.
+5. **Selective accelerator backends** for the workloads that
+   benefit from them. Neural causal estimators (TARNet, CFRNet,
+   DragonNet, CEVAE, DeepIV) route through PyTorch CUDA / MPS via
+   the `STATSPAI_TORCH_DEVICE` environment variable. The HDFE
+   residualiser exposes `backend="jax"` and `sp.fast.feols_jax`
+   runs the full WLS solve on JAX / XLA. The largest GPU win in
+   v1.14 is `sp.fast.feols_jax_bootstrap`: the same JIT-compiled
+   WLS kernel is lifted to a `jax.vmap` batched primitive, giving
+   a 10--100x speedup over sequential CPU bootstrap on CUDA / TPU
+   at $B \geq 1{,}000$. Four bootstrap variants share the same
+   JAX kernel infrastructure --- pairs, cluster, wild, and wild
+   cluster [@cameron2008bootstrap] --- with the wild
+   variants using the score formulation $\hat\beta^*_b =
+   \hat\beta + (X'WX)^{-1}\, X'W\, (\eta_b \odot \hat u)$ which
+   is mathematically identical to refitting on
+   $y^* = X\hat\beta + \eta \odot \hat u$ but needs only a single
+   matrix--vector multiply per iteration. The remainder of the
+   package (DiD, RD, synthetic control, GMM) is CPU-only by
+   design: these workloads are bandwidth-bound or involve small-$K$
+   convex programs where a tuned Rust + Numba kernel matches GPU
+   performance. Cluster-robust meat matrices use a Rust + Rayon
+   kernel parallel over clusters
+   (`statspai.core._numba_kernels.cluster_meat`) introduced in
+   v1.14, and `sp.iv(absorb=...)` newly wires the Rust HDFE
+   residualiser into the 2SLS first stage.
 
 The package is implemented in pure Python atop NumPy, SciPy, Pandas,
 statsmodels, scikit-learn, and linearmodels, with optional PyTorch and
-JAX backends. It supports Python $\geq$ 3.9 and is distributed via
-PyPI under the MIT license.
+JAX backends and a Rust HDFE / cluster-meat kernel
+(`statspai_hdfe`, PyO3 + Rayon) that the Python wrappers prefer when
+built and silently fall through when absent. It supports Python
+$\geq$ 3.9 and is distributed via PyPI under the MIT license.
 
 # Validation
 

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "StatsPAI"
-version = "1.13.1"
+version = "1.14.0"
 description = "The Agent-Native Causal Inference & Econometrics Toolkit for Python"
 readme = "README.md"
 license = {text = "MIT"}

src/statspai/__init__.py

@@ -22,7 +22,7 @@
 >>> sp.outreg2(result, filename="results.xlsx")
 """
 
-__version__ = "1.13.1"
+__version__ = "1.14.0"
 __author__ = "Biaoyue Wang"
 __email__ = "brycew6m@stanford.edu"