CI: run optional backend tests on pull requests

.github/workflows/optional-backend-tests.yml

@@ -0,0 +1,69 @@
+name: Optional backend tests
+
+on:
+  pull_request:
+  push:
+    branches:
+      - master
+  workflow_dispatch:
+
+permissions:
+  contents: read
+
+concurrency:
+  group: optional-backends-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  test:
+    name: ${{ matrix.backend }} (Python ${{ matrix.python-version }})
+    runs-on: ubuntu-latest
+    timeout-minutes: 30
+
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.11", "3.12"]
+        backend: [tensorflow, pytorch, jax, jax-pytorch-parity]
+        include:
+          - backend: tensorflow
+            extras: tf
+            test-path: tests/
+            marker: tf
+            pytest-flags: --runtf
+          - backend: pytorch
+            extras: torch
+            test-path: tests/
+            marker: torch and not jax
+            pytest-flags: --runtorch
+          - backend: jax
+            extras: jax
+            test-path: tests/
+            marker: jax and not torch
+            pytest-flags: --runjax
+          - backend: jax-pytorch-parity
+            extras: torch,jax
+            test-path: tests/
+            marker: jax and torch
+            pytest-flags: --runjax --runtorch
+
+    steps:
+      - uses: actions/checkout@v6
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v6
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Upgrade packaging tools
+        run: python -m pip install --upgrade pip setuptools
+      - name: Install CPU-only PyTorch
+        if: contains(matrix.extras, 'torch')
+        run: python -m pip install --index-url https://download.pytorch.org/whl/cpu torch
+      - name: Install dependencies
+        run: python -m pip install -e ".[test,${{ matrix.extras }}]"
+      - name: Run ${{ matrix.backend }} tests
+        env:
+          OMP_NUM_THREADS: "1"
+        run: >-
+          python -m pytest -vs ${{ matrix.test-path }}
+          -m "${{ matrix.marker }}"
+          ${{ matrix.pytest-flags }}

tests/test_cevae.py

@@ -1,3 +1,4 @@
+import numpy as np
 import pandas as pd
 import pytest
 
@@ -9,20 +10,25 @@
 from causalml.dataset import simulate_hidden_confounder
 from causalml.metrics import get_cumgain
 
+from .const import RANDOM_SEED
+
 
 @pytest.mark.torch
 def test_CEVAE():
-    y, X, treatment, tau, b, e = simulate_hidden_confounder(
-        n=10000, p=5, sigma=1.0, adj=0.0
+    np.random.seed(RANDOM_SEED)
+    torch.manual_seed(RANDOM_SEED)
+    y, X, treatment, tau, _, _ = simulate_hidden_confounder(
+        n=2000, p=5, sigma=1.0, adj=0.0
     )
 
     outcome_dist = "normal"
     latent_dim = 20
-    hidden_dim = 200
-    num_epochs = 50
+    hidden_dim = 64
+    num_epochs = 20
     batch_size = 100
     learning_rate = 1e-3
     learning_rate_decay = 0.1
+    num_samples = 200
 
     cevae = CEVAE(
         outcome_dist=outcome_dist,
@@ -32,6 +38,7 @@ def test_CEVAE():
         batch_size=batch_size,
         learning_rate=learning_rate,
         learning_rate_decay=learning_rate_decay,
+        num_samples=num_samples,
     )
 
     cevae.fit(
@@ -40,17 +47,16 @@ def test_CEVAE():
         y=torch.tensor(y, dtype=torch.float),
     )
 
-    # check the accuracy of the ite accuracy
     ite = cevae.predict(X).flatten()
 
-    auuc_metrics = pd.DataFrame(
-        {"ite": ite, "W": treatment, "y": y, "treatment_effect_col": tau}
-    )
+    auuc_metrics = pd.DataFrame({"ite": ite, "W": treatment, "y": y, "tau": tau})
 
     cumgain = get_cumgain(
         auuc_metrics, outcome_col="y", treatment_col="W", treatment_effect_col="tau"
     )
 
-    # Check if the cumulative gain when using the model's prediction is
-    # higher than it would be under random targeting
-    assert cumgain["ite"].sum() > cumgain["Random"].sum()
+    # Compare the model's cumulative gain with random targeting.
+    random_gain = np.linspace(
+        cumgain["ite"].iloc[0], cumgain["ite"].iloc[-1], cumgain.shape[0]
+    )
+    assert cumgain["ite"].sum() > random_gain.sum()