feat: add AIR case study and harden ADEV float0 handling

- port a GenJAX-only AIR example (core/main/figs) with docs, tasks, and tests
- align ADEV primitive JVP dispatch with JAX float0->Zero canonicalization
- add staged flip_mvd regression coverage for seed+vmap+jit bool-cast paths
- document AD zero-tangent canonicalization behavior in ADEV guides

Author: femtomc

README.md

@@ -339,6 +339,31 @@ pixi run -e faircoin python -m examples.faircoin.main \
 
 **Outputs**: `figs/faircoin_combined_posterior_and_timing_obs50_samples2000.pdf`
 
+### AIR Estimator Study (GenJAX-only)
+
+**What it does**: Trains an AIR-style latent-variable model and compares GenJAX discrete-gradient estimators (`enum`, `reinforce`, `mvd`, `hybrid`) under a shared objective/training loop.
+
+**Figures in the paper**: Port of the PLDI'24 AIR estimator experiment (GenJAX path only).
+
+**Commands**:
+```bash
+# quick smoke comparison (small architecture + synthetic prior samples)
+pixi run air-compare
+
+# single-estimator run with CSV outputs
+pixi run python -m examples.air.main train \
+  --estimator enum \
+  --small-config \
+  --num-examples 512 \
+  --epochs 6 \
+  --history-csv figs/air_enum_history.csv \
+  --summary-csv figs/air_enum_summary.csv
+```
+
+**Dataset modes**:
+- `--dataset synthetic` (default): samples from the AIR prior (no extra framework dependencies).
+- `--dataset multi-mnist --data-path /path/to/multi_mnist_uint8.npz`: load pre-generated multi-MNIST arrays.
+
 ### Curve Fitting with Outlier Detection
 
 **What it does**: Polynomial regression with robust outlier detection, demonstrating:
@@ -477,6 +502,10 @@ Running without CUDA executes the same probabilistic program, but the SMC benchm
 
 All figures are saved to `figs/`:
 
+### AIR
+- `air_*_history.csv` - Per-epoch objective/accuracy/time logs (optional)
+- `air_*_summary.csv` - Estimator summary table (optional)
+
 ### Faircoin
 - `faircoin_combined_posterior_and_timing_obs50_samples2000.pdf` - Framework comparison (timing + posterior accuracy)
 

examples/AGENTS.md

@@ -5,6 +5,7 @@ Each case has its own `AGENTS.md`—read it before editing that case.
 
 ## Case Study Map
 
+- `air/`: AIR estimator comparison (GenJAX-only port from PLDI'24 artifact)
 - `cone/`: ADEV objective variants (ELBO / IWAE / HVI-style families)
 - `curvefit/`: polynomial regression + scalable inference + outlier modeling
 - `faircoin/`: Beta–Bernoulli baseline benchmarking

examples/air/AGENTS.md

@@ -0,0 +1,48 @@
+# AIR Case Study Guide
+
+This case ports the GenJAX AIR estimator experiment from the PLDI'24 programmable-VI artifact.
+
+## Purpose
+
+Evaluate discrete-gradient estimators in an AIR-style latent-variable model:
+
+- `enum`
+- `reinforce`
+- `mvd`
+- `hybrid` (MVD for early `z_pres` sites, ENUM on the final site)
+
+## Key Files
+
+- `core.py`: model/guide definitions, STN transforms, objectives, training/eval utilities
+- `main.py`: CLI (`train`, `compare`)
+
+## CLI Commands
+
+```bash
+# Quick smoke comparison on synthetic data
+pixi run python -m examples.air.main compare --small-config --num-examples 128 --epochs 2
+
+# Train a single estimator and save history
+pixi run python -m examples.air.main train \
+  --estimator enum \
+  --small-config \
+  --num-examples 256 \
+  --epochs 4 \
+  --history-csv figs/air_enum_history.csv
+
+# Use pre-generated multi-MNIST data (multi_mnist_uint8.npz)
+pixi run python -m examples.air.main compare \
+  --dataset multi-mnist \
+  --data-path /path/to/multi_mnist_uint8.npz \
+  --num-examples 2000
+```
+
+## Notes
+
+- Default dataset mode is `synthetic` (samples from the AIR prior), so the case runs without Pyro/Torch.
+- For `multi-mnist`, supply an existing NPZ file.
+- Keep heavy logic in `core.py`; keep `main.py` as orchestration only.
+
+## Tests
+
+- `tests/test_air_example.py`

examples/air/__init__.py

@@ -0,0 +1 @@
+"""AIR case study (GenJAX-only port of the PLDI'24 artifact experiment)."""