{
"cells": [
{
"cell_type": "markdown",
"id": "88e2895d",
"metadata": {},
"source": [
"# CEVAE: JAX vs PyTorch Backend Benchmark\n",
"\n",
"This notebook compares the two CEVAE (Causal Effect Variational Autoencoder,\n",
"Louizos et al. 2017) backends available in CausalML:\n",
"\n",
"- `causalml.inference.torch.CEVAE` — the original PyTorch/[pyro](https://pyro.ai) implementation\n",
" (wraps `pyro.contrib.cevae`)\n",
"- `causalml.inference.jax.CEVAE` — the JAX/[`flax.nnx`](https://flax.readthedocs.io/en/latest/nnx/index.html)\n",
" implementation (see `causalml/inference/jax/cevae/`)\n",
"\n",
"Both models are trained with **identical hyperparameters** on the same synthetic\n",
"dataset with a hidden confounder, and compared on ATE, MAE (against the true ITE),\n",
"AUUC, and wall-clock training time.\n",
"\n",
"Requires the optional dependencies: `pip install \"causalml[jax,torch]\"`."
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "a0f92188",
"metadata": {
"execution": {
"iopub.execute_input": "2026-07-18T15:18:32.042684Z",
"iopub.status.busy": "2026-07-18T15:18:32.042567Z",
"iopub.status.idle": "2026-07-18T15:18:35.222393Z",
"shell.execute_reply": "2026-07-18T15:18:35.222011Z"
}
},
"outputs": [],
"source": [
"# Avoid OMP/threading conflicts when torch and JAX share a single process on macOS.\n",
"import logging\n",
"import os\n",
"import warnings\n",
"\n",
"os.environ.setdefault(\"KMP_DUPLICATE_LIB_OK\", \"TRUE\")\n",
"os.environ.setdefault(\"OMP_NUM_THREADS\", \"1\")\n",
"\n",
"warnings.filterwarnings(\"ignore\")\n",
"\n",
"import time\n",
"\n",
"import numpy as np\n",
"import pandas as pd\n",
"import torch\n",
"import pyro\n",
"from matplotlib import pyplot as plt\n",
"import seaborn as sns\n",
"\n",
"from sklearn.model_selection import train_test_split\n",
"from sklearn.metrics import mean_absolute_error\n",
"\n",
"from causalml.dataset import simulate_hidden_confounder\n",
"from causalml.inference.torch import CEVAE as CEVAETorch\n",
"from causalml.inference.jax import CEVAE as CEVAEJax\n",
"from causalml.metrics import auuc_score, plot_gain\n",
"\n",
"# causalml.inference.torch enables DEBUG logging for pyro at import time; quiet it.\n",
"logging.getLogger(\"pyro\").setLevel(logging.WARNING)\n",
"\n",
"plt.style.use(\"fivethirtyeight\")\n",
"sns.set_palette(\"Paired\")\n",
"plt.rcParams[\"figure.figsize\"] = (12, 8)"
]
},
{
"cell_type": "markdown",
"id": "73442d0e",
"metadata": {},
"source": [
"## Synthetic data with a hidden confounder\n",
"\n",
"`simulate_hidden_confounder` generates data where a latent variable biases both\n",
"treatment assignment and outcome — the setting CEVAE was designed for. The\n",
"ground-truth individual treatment effect `tau` is known, so we can score each\n",
"backend directly against it."
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "b80133e8",
"metadata": {
"execution": {
"iopub.execute_input": "2026-07-18T15:18:35.224144Z",
"iopub.status.busy": "2026-07-18T15:18:35.224024Z",
"iopub.status.idle": "2026-07-18T15:18:35.229477Z",
"shell.execute_reply": "2026-07-18T15:18:35.229154Z"
}
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"train: 16000 rows | validation: 4000 rows\n",
"true ATE (train): 0.7293 | true ATE (val): 0.7234\n"
]
}
],
"source": [
"np.random.seed(123)\n",
"\n",
"y, X, w, tau, b, e = simulate_hidden_confounder(n=20000, p=5, sigma=1.0, adj=0.0)\n",
"\n",
"(\n",
" X_train,\n",
" X_val,\n",
" y_train,\n",
" y_val,\n",
" w_train,\n",
" w_val,\n",
" tau_train,\n",
" tau_val,\n",
") = train_test_split(X, y, w, tau, test_size=0.2, random_state=123, shuffle=True)\n",
"\n",
"print(f\"train: {X_train.shape[0]} rows | validation: {X_val.shape[0]} rows\")\n",
"print(f\"true ATE (train): {tau_train.mean():.4f} | true ATE (val): {tau_val.mean():.4f}\")"
]
},
{
"cell_type": "markdown",
"id": "3893a17b",
"metadata": {},
"source": [
"## Shared hyperparameters\n",
"\n",
"Both backends expose the same sklearn-style API and are configured identically,\n",
"so any difference in the results comes from the implementations themselves."
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "661ce623",
"metadata": {
"execution": {
"iopub.execute_input": "2026-07-18T15:18:35.230522Z",
"iopub.status.busy": "2026-07-18T15:18:35.230456Z",
"iopub.status.idle": "2026-07-18T15:18:35.232045Z",
"shell.execute_reply": "2026-07-18T15:18:35.231710Z"
}
},
"outputs": [],
"source": [
"common = dict(\n",
" outcome_dist=\"normal\",\n",
" latent_dim=20,\n",
" hidden_dim=200,\n",
" num_epochs=25,\n",
" batch_size=1000,\n",
" learning_rate=1e-3,\n",
" learning_rate_decay=0.1,\n",
" num_layers=3,\n",
" num_samples=100,\n",
")"
]
},
{
"cell_type": "markdown",
"id": "1863a3e0",
"metadata": {},
"source": [
"## PyTorch (pyro) backend"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "fe55106a",
"metadata": {
"execution": {
"iopub.execute_input": "2026-07-18T15:18:35.232977Z",
"iopub.status.busy": "2026-07-18T15:18:35.232925Z",
"iopub.status.idle": "2026-07-18T15:19:00.049983Z",
"shell.execute_reply": "2026-07-18T15:19:00.049325Z"
}
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"PyTorch backend trained in 6.7s\n"
]
}
],
"source": [
"pyro.set_rng_seed(0)\n",
"torch.manual_seed(0)\n",
"\n",
"torch_cevae = CEVAETorch(**common)\n",
"\n",
"t0 = time.perf_counter()\n",
"torch_cevae.fit(\n",
" X=torch.tensor(X_train, dtype=torch.float),\n",
" treatment=torch.tensor(w_train, dtype=torch.float),\n",
" y=torch.tensor(y_train, dtype=torch.float),\n",
")\n",
"torch_time = time.perf_counter() - t0\n",
"\n",
"ite_torch_train = torch_cevae.predict(X_train).ravel()\n",
"ite_torch_val = torch_cevae.predict(X_val).ravel()\n",
"\n",
"print(f\"PyTorch backend trained in {torch_time:.1f}s\")"
]
},
{
"cell_type": "markdown",
"id": "f2751182",
"metadata": {},
"source": [
"## JAX (flax.nnx) backend"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "807d99c4",
"metadata": {
"execution": {
"iopub.execute_input": "2026-07-18T15:19:00.052621Z",
"iopub.status.busy": "2026-07-18T15:19:00.051903Z",
"iopub.status.idle": "2026-07-18T15:19:12.518640Z",
"shell.execute_reply": "2026-07-18T15:19:12.518079Z"
}
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"JAX backend trained in 4.6s\n"
]
}
],
"source": [
"jax_cevae = CEVAEJax(**common, seed=0)\n",
"\n",
"t0 = time.perf_counter()\n",
"jax_cevae.fit(X_train, w_train, y_train)\n",
"jax_time = time.perf_counter() - t0\n",
"\n",
"ite_jax_train = jax_cevae.predict(X_train).ravel()\n",
"ite_jax_val = jax_cevae.predict(X_val).ravel()\n",
"\n",
"print(f\"JAX backend trained in {jax_time:.1f}s\")"
]
},
{
"cell_type": "markdown",
"id": "9e8eddb3",
"metadata": {},
"source": [
"## Benchmark results\n",
"\n",
"### Training set"
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "2a4c74c0",
"metadata": {
"execution": {
"iopub.execute_input": "2026-07-18T15:19:12.520488Z",
"iopub.status.busy": "2026-07-18T15:19:12.520247Z",
"iopub.status.idle": "2026-07-18T15:19:12.598030Z",
"shell.execute_reply": "2026-07-18T15:19:12.597621Z"
}
},
"outputs": [
{
"data": {
"text/html": [
"\n",
"\n",
"
\n",
" \n",
" \n",
" | \n",
" ATE | \n",
" MAE | \n",
" fit time (s) | \n",
" AUUC | \n",
"
\n",
" \n",
" \n",
" \n",
" | CEVAE (PyTorch) | \n",
" 0.884192 | \n",
" 0.210640 | \n",
" 6.747957 | \n",
" 0.524815 | \n",
"
\n",
" \n",
" | CEVAE (JAX) | \n",
" 0.831832 | \n",
" 0.205091 | \n",
" 4.628406 | \n",
" 0.545208 | \n",
"
\n",
" \n",
" | actual | \n",
" 0.729305 | \n",
" NaN | \n",
" NaN | \n",
" NaN | \n",
"
\n",
" \n",
"
\n",
"
"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"plot_gain(df_preds_train)"
]
},
{
"cell_type": "markdown",
"id": "78788b42",
"metadata": {},
"source": [
"### Validation set"
]
},
{
"cell_type": "code",
"execution_count": 8,
"id": "19ee1bb8",
"metadata": {
"execution": {
"iopub.execute_input": "2026-07-18T15:19:12.693524Z",
"iopub.status.busy": "2026-07-18T15:19:12.693440Z",
"iopub.status.idle": "2026-07-18T15:19:12.713625Z",
"shell.execute_reply": "2026-07-18T15:19:12.713257Z"
}
},
"outputs": [
{
"data": {
"text/html": [
"\n",
"\n",
"
\n",
" \n",
" \n",
" | \n",
" ATE | \n",
" MAE | \n",
" AUUC | \n",
"
\n",
" \n",
" \n",
" \n",
" | CEVAE (PyTorch) | \n",
" 0.884965 | \n",
" 0.215583 | \n",
" 0.524284 | \n",
"
\n",
" \n",
" | CEVAE (JAX) | \n",
" 0.831255 | \n",
" 0.208282 | \n",
" 0.545295 | \n",
"
\n",
" \n",
" | actual | \n",
" 0.723394 | \n",
" NaN | \n",
" NaN | \n",
"
\n",
" \n",
"
\n",
"