Add Python package structure, simplify API

.github/workflows/CI.yml → .github/workflows/CI-Julia.yml

@@ -1,4 +1,4 @@
-name: CI
+name: CI-Julia
 on:
   push:
     branches:
@@ -7,16 +7,14 @@ on:
   pull_request:
   workflow_dispatch:
 concurrency:
-  # Skip intermediate builds: always.
-  # Cancel intermediate builds: only if it is a pull request build.
   group: ${{ github.workflow }}-${{ github.ref }}
   cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }}
 jobs:
   test:
     name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }}
     runs-on: ${{ matrix.os }}
     timeout-minutes: 60
-    permissions: # needed to allow julia-actions/cache to proactively delete old caches that it has created
+    permissions:
       actions: write
       contents: read
     strategy:
@@ -36,7 +34,6 @@ jobs:
           arch: ${{ matrix.arch }}
       - uses: julia-actions/cache@v2
       - name: Install dependencies & run tests
-        # how to add the local DIT to the DI test env?
         run: julia --project=./julia --color=yes -e '
           using Pkg;
           Pkg.test("SmoothedDifferentiation"; coverage=true);'
@@ -53,10 +50,6 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v5
-      # - uses: julia-actions/setup-julia@v2
-      #   with:
-      #     version: '1'
-      # - uses: julia-actions/cache@v2
       - uses: fredrikekre/runic-action@v1
         with:
-          version: '1'
+          version: '1'

.github/workflows/CI-Python.yml

@@ -0,0 +1,48 @@
+name: CI-Python
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+  workflow_dispatch:
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }}
+jobs:
+  test:
+    name: Python ${{ matrix.python-version }}
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.12", "3.14"]
+    steps:
+      - uses: actions/checkout@v5
+      - name: Install uv
+        uses: astral-sh/setup-uv@v5
+      - name: Set up Python
+        run: uv python install ${{ matrix.python-version }}
+      - name: Install dependencies
+        run: uv sync --group test
+        working-directory: python
+      - name: Test
+        run: uv run --group test pytest
+        working-directory: python
+  lint:
+    name: Ruff
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v5
+      - name: Install uv
+        uses: astral-sh/setup-uv@v5
+      - name: Set up Python
+        run: uv python install 3.14
+      - name: Install dependencies
+        run: uv sync --group test
+        working-directory: python
+      - name: Lint
+        run: uv run ruff check .
+        working-directory: python
+      - name: Format
+        run: uv run ruff format --check .
+        working-directory: python

.gitignore

@@ -58,6 +58,7 @@ coverage.xml
 .hypothesis/
 .pytest_cache/
 cover/
+*.lock
 
 # Translations
 *.mo

README.md

@@ -22,6 +22,10 @@ Code for the NeurIPS 2025 paper *"Smoothed Differentiation Efficiently Mitigates
 
 ## Experiments
 
+> [!NOTE]
+> The experiment code refers to the state of the repository at the time of publication.
+> See the [`neurips25`](https://github.com/adrhill/smoothdiff-experiments/tree/neurips25) tag for the exact code used in the paper.
+
 We provide all code and virtual environments required to reproduce our experiments and figures.
 
 ### Running experiments

python/README.md

@@ -0,0 +1,35 @@
+# smoothdiff-torch
+
+SmoothDiff implementation in PyTorch.
+
+## Installation
+
+```bash
+pip install "smoothdiff-torch @ git+https://github.com/adrhill/smoothdiff-experiments.git#subdirectory=python"
+```
+
+## Usage
+
+```python
+import smoothdiff_torch as smoothdiff
+
+# Prepare model by replacing non-linear layers
+model = smoothdiff.prepare_model(model)
+
+# Collect gradient statistics over multiple noisy forward passes
+smoothdiff.set_mode(model, collect_stats=True, smooth_backward=False)
+with torch.no_grad():
+    for _ in range(n_samples):
+        model(inputs + torch.randn_like(inputs) * std)
+
+# Compute SmoothDiff explanation via a single backward pass
+smoothdiff.set_mode(model, collect_stats=False, smooth_backward=True)
+inputs = inputs.clone().detach().requires_grad_(True)
+model(inputs).max().backward()
+attribution = inputs.grad
+
+# Free accumulated statistics
+smoothdiff.reset_stats(model)
+```
+
+See [`examples/smoothdiff.ipynb`](examples/smoothdiff.ipynb) for a full example.

python/examples/smoothdiff.ipynb

@@ -0,0 +1,207 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "sys.path.insert(0, \"..\")\n",
+    "\n",
+    "import numpy as np\n",
+    "from PIL import Image\n",
+    "import cmcrameri.cm as cmc\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "import torch\n",
+    "import torch.nn.functional as F\n",
+    "import torchvision\n",
+    "from torchvision import transforms\n",
+    "\n",
+    "import smoothdiff_torch as smoothdiff\n",
+    "\n",
+    "DEVICE = \"mps\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Load Model\n",
+    "model = torchvision.models.vgg16(weights=\"IMAGENET1K_V1\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Load input image\n",
+    "img = Image.open(\"assets/car.jpg\")\n",
+    "img = img.convert(\"RGB\")\n",
+    "img"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Apply preprocessing to image\n",
+    "preprocess = transforms.Compose(\n",
+    "    [\n",
+    "        transforms.Resize(256),\n",
+    "        transforms.CenterCrop(224),\n",
+    "        transforms.ToTensor(),\n",
+    "        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n",
+    "    ]\n",
+    ")\n",
+    "\n",
+    "inputs = preprocess(img).unsqueeze(0).to(DEVICE)\n",
+    "inputs.shape"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Run SmoothDiff\n",
+    "def smoothdiff_explanation(model, inputs, n_samples=20, std=0.5):\n",
+    "    # Replacing non-linear functions by ones that compute VEJPs on the backward-pass\n",
+    "    model = smoothdiff.prepare_model(model)\n",
+    "    model.to(DEVICE).eval()\n",
+    "\n",
+    "    # Step 1: Compute n forward passes\n",
+    "    smoothdiff.set_mode(model, collect_stats=True, smooth_backward=False)\n",
+    "    with torch.no_grad():\n",
+    "        for _ in range(n_samples):\n",
+    "            model(inputs + torch.randn_like(inputs) * std)\n",
+    "\n",
+    "    # Step 2: Compute single SmoothDiff backward-pass (VEJPs)\n",
+    "    smoothdiff.set_mode(model, collect_stats=False, smooth_backward=True)\n",
+    "    inputs = inputs.clone().detach().requires_grad_(True).to(DEVICE)  # Y'all should really give Julia a shot!\n",
+    "    y = model(inputs)\n",
+    "    y.max().backward()  # explain w.r.t. maximally activated output neuron\n",
+    "\n",
+    "    attribution = inputs.grad.clone().cpu().detach().numpy()\n",
+    "    return attribution"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# In case you are curious about the layers SmoothDiff replaces:\n",
+    "smoothdiff.prepare_model(model)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Heatmapping utilities\n",
+    "def pool(attribution: np.ndarray) -> np.ndarray:\n",
+    "    \"\"\"\n",
+    "    Pool attributions by computing the norm over color channels\n",
+    "    (summation should not be used for sensitivity-based methods).\n",
+    "    \"\"\"\n",
+    "    assert attribution.shape == (3, 224, 224), (\n",
+    "        f\"Shape mismatch: expected (3, 224, 224), got {attribution.shape}\"\n",
+    "    )\n",
+    "    return np.linalg.norm(attribution, axis=0)\n",
+    "\n",
+    "\n",
+    "def normalize(attribution: np.ndarray, percentile=99.9) -> np.ndarray:\n",
+    "    \"\"\"\n",
+    "    Normalize by max-abs. value of attribution, removing outliers by computing the 99.9th percentile of values.\n",
+    "    \"\"\"\n",
+    "    absmax = np.percentile(np.abs(attribution), percentile)\n",
+    "    return np.clip(attribution / absmax, a_min=-1.0, a_max=1.0)\n",
+    "\n",
+    "\n",
+    "def heatmap(ax, attribution):\n",
+    "    return ax.imshow(normalize(pool(attribution)), cmap=cmc.batlow, vmin=0.0, vmax=1.0)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Plot SmoothDiff explanations over increasing sample sizes\n",
+    "std = 0.5\n",
+    "ns = [2, 4, 8, 16, 32, 64]\n",
+    "\n",
+    "fig, axes = plt.subplots(ncols=len(ns), figsize=(18, 3))\n",
+    "for i, n_samples in enumerate(ns):\n",
+    "    ax = axes[i]\n",
+    "    attribution = smoothdiff_explanation(model, inputs, n_samples, std)\n",
+    "    smoothdiff.reset_stats(model) # reset VEJP statistics\n",
+    "\n",
+    "    heatmap(ax, attribution[0])\n",
+    "    ax.title.set_text(f\"n = {n_samples}\")\n",
+    "    ax.axis(\"off\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Cleanup\n",
+    "SmoothDiff layers store accumulated gradient statistics from the forward passes. Call `reset_smoothdiff_stats` to free this memory after you're done computing explanations."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "smoothdiff.reset_stats(model)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "SmoothDiff (PyTorch)",
+   "language": "python",
+   "name": "smoothdiff-torch"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.8"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}