Add JAX/flax.nnx backend for CEVAE

[xrhd]

Proposed changes

Adds a JAX backend for CEVAE (Causal Effect Variational Autoencoder, Louizos et al. 2017) under causalml/inference/jax/cevae/, using flax.nnx as the neural network framework and optax for optimization. Exposes the same sklearn-style API as the existing PyTorch implementation (fit, predict, fit_predict, save, load) and follows the opt-in causalml[jax] extra pattern established by #918. A benchmark notebook comparing the two backends side-by-side is included.

Resolves #956

Types of changes

• New feature (non-breaking change which adds functionality)

Checklist

• I have read the CONTRIBUTING doc
• I have signed the CLA
• Lint and unit tests pass locally with my changes
• I have added tests that prove my fix is effective or that my feature works
• I have added necessary documentation (if appropriate)

Changes

• causalml/inference/jax/cevae/modeling.py — flax.nnx modules porting pyro.contrib.cevae: FullyConnected, outcome nets (Bernoulli/Exponential/Laplace/Normal/StudentT), DiagNormalNet, PreWhitener, Model, Guide
• causalml/inference/jax/cevae/losses.py — ELBO loss (cevae_loss, @jax.jit-compiled)
• causalml/inference/jax/cevae/cevae.py — sklearn-style CEVAE wrapper; training via optax.adamw with exponential LR decay matching pyro's ClippedAdam schedule; serialization via orbax-checkpoint; hardened with a non-finite-loss guard during training and feature-dimension validation at predict time
• causalml/inference/jax/cevae/README.md — module documentation
• causalml/inference/jax/__init__.py — export CEVAE
• tests/test_jax_cevae.py — fit/predict shapes (exercising the default studentt outcome distribution), ATE recovery, save/load round-trip, and a JAX↔Pyro parity test (each backend runs in its own subprocess to avoid OMP/threading conflicts)
• docs/examples/cevae_jax_vs_torch.ipynb — benchmark notebook comparing the JAX and PyTorch backends on synthetic data with identical hyperparameters (ATE, MAE, AUUC, training time)

Tests

pip install ".[jax,test]"
pytest tests/test_jax_cevae.py --runjax -v
# parity test additionally requires torch/pyro:
pytest tests/test_jax_cevae.py --runjax --runtorch -v

All 4 tests pass locally (macOS arm64, Python 3.13): shapes, ATE recovery, save/load round-trip, and JAX↔Pyro parity.

Dependencies

No dependency changes — uses the existing opt-in [jax] extra (jax, flax, optax, orbax-checkpoint).

Further comments

The implementation mirrors pyro.contrib.cevae module/parameter names so the two backends can be diffed line-by-line. Training is a pure JAX loop with nnx.jit. The dependency stays opt-in via causalml[jax].

The second commit is a polish pass from self-review: ceil-division fix in the LR schedule step count, NaN guard during training, predict-time feature-dimension validation, test coverage for the default studentt outcome distribution, plus the JAX vs PyTorch benchmark notebook.

References

• Louizos et al. (2017), "Causal Effect Inference with Deep Latent-Variable Models" — http://papers.nips.cc/paper/7223-causal-effect-inference-with-deep-latent-variable-models.pdf
• Precedent: Add JAX/flax.nnx backend for DragonNet #917 / Add JAX/flax.nnx backend for DragonNet #918 (JAX/flax.nnx backend for DragonNet)
• flax.nnx: https://flax.readthedocs.io/en/latest/nnx/index.html
• JAX ecosystem: https://docs.jaxstack.ai/en/latest/index.html

[jeongyoonlee]

LGTM — thanks for the contribution, and welcome!