Align the jax stack, name the CUDA extra for CUDA, and drop Python 3.11

.github/workflows/build-verification.yml

@@ -15,7 +15,7 @@ jobs:
       fail-fast: false
       matrix:
         os: [ubuntu-latest, macos-14]
-        python-version: ["3.11", "3.12"]
+        python-version: ["3.12", "3.13"]
 
     steps:
       - uses: actions/checkout@v4

.github/workflows/ci.yml

@@ -25,7 +25,7 @@ jobs:
       - name: Set up DiffBio
         uses: ./.github/actions/setup-diffbio
         with:
-          python-version: "3.11"
+          python-version: "3.12"
           cache-suffix: quality
           linux-editable-target: ".[dev]"
 
@@ -120,7 +120,7 @@ jobs:
       - name: Set up DiffBio
         uses: ./.github/actions/setup-diffbio
         with:
-          python-version: "3.11"
+          python-version: "3.12"
           cache-suffix: unit-${{ matrix.shard.id }}
           # Install genomics extra so pyfaidx / pysam are present for the
           # error-path tests in tests/sources/test_{fasta,bam}.py.
@@ -180,7 +180,7 @@ jobs:
       - name: Set up DiffBio
         uses: ./.github/actions/setup-diffbio
         with:
-          python-version: "3.11"
+          python-version: "3.12"
           cache-suffix: integration
           linux-editable-target: ".[dev,test]"
 
@@ -226,7 +226,7 @@ jobs:
       - name: Set up DiffBio
         uses: ./.github/actions/setup-diffbio
         with:
-          python-version: "3.11"
+          python-version: "3.12"
           cache-suffix: e2e
           linux-editable-target: ".[dev,test]"
 
@@ -287,7 +287,7 @@ jobs:
       - name: Set up DiffBio
         uses: ./.github/actions/setup-diffbio
         with:
-          python-version: "3.11"
+          python-version: "3.12"
           cache-suffix: perf
           # chem carries rdkit, needed by tests/benchmarks/test_bench_dti.py and
           # test_bench_molnet.py.
@@ -338,7 +338,7 @@ jobs:
       - name: Set up DiffBio
         uses: ./.github/actions/setup-diffbio
         with:
-          python-version: "3.11"
+          python-version: "3.12"
           cache-suffix: coverage
           # chem carries rdkit, which the drug-discovery, molnet and molecular-splitter
           # tests exercise. It is an extra so that installing the package does not pull a

.github/workflows/docs.yml

@@ -27,8 +27,8 @@ jobs:
       - name: Set up DiffBio docs environment
         uses: ./.github/actions/setup-diffbio
         with:
-          python-version: '3.11'
-          cache-suffix: docs-py311
+          python-version: '3.12'
+          cache-suffix: docs-py312
           linux-editable-target: ".[docs]"
 
       - name: Build documentation

.github/workflows/publish.yml

@@ -41,7 +41,7 @@ jobs:
       - name: Set up Python
         uses: actions/setup-python@v6
         with:
-          python-version: "3.11"
+          python-version: "3.12"
 
       - name: Install build dependencies
         run: |

.github/workflows/quality-checks.yml

@@ -14,7 +14,7 @@ jobs:
       - name: Set up DiffBio
         uses: ./.github/actions/setup-diffbio
         with:
-          python-version: "3.11"
+          python-version: "3.12"
           cache-suffix: quality-report
           linux-editable-target: ".[dev]"
 

.github/workflows/security.yml

@@ -20,7 +20,7 @@ jobs:
       - name: Set up DiffBio
         uses: ./.github/actions/setup-diffbio
         with:
-          python-version: "3.11"
+          python-version: "3.12"
           cache-suffix: security
           linux-editable-target: ".[dev]"
 

.readthedocs.yaml

@@ -6,7 +6,7 @@ version: 2
 build:
   os: ubuntu-22.04
   tools:
-    python: "3.11"
+    python: "3.12"
   jobs:
     post_install:
       - pip install mkdocs-material

CONTRIBUTING.md

@@ -20,7 +20,7 @@ source ./activate.sh
 
 ### Prerequisites
 
-- Python 3.11+
+- Python 3.12+
 - uv package manager
 - Git
 

README.md

@@ -1,7 +1,7 @@
 # DiffBio
 
 <p align="center">
-  <a href="https://www.python.org/downloads/"><img src="https://img.shields.io/badge/python-3.11+-blue.svg" alt="Python 3.11+"></a>
+  <a href="https://www.python.org/downloads/"><img src="https://img.shields.io/badge/python-3.12+-blue.svg" alt="Python 3.12+"></a>
   <a href="https://jax.readthedocs.io/"><img src="https://img.shields.io/badge/JAX-0.6.1+-green.svg" alt="JAX"></a>
   <a href="https://flax.readthedocs.io/"><img src="https://img.shields.io/badge/Flax-0.12+-orange.svg" alt="Flax"></a>
   <a href="LICENSE"><img src="https://img.shields.io/badge/license-MIT-blue.svg" alt="License"></a>
@@ -293,7 +293,7 @@ DiffBio/
 
 ## Requirements
 
-- Python 3.11+
+- Python 3.12+
 - JAX 0.6.1+
 - Flax 0.12+
 - Optax 0.1.4+

benchmarks/_optimizers.py

@@ -16,12 +16,19 @@ def create_benchmark_optimizer(
     learning_rate: float,
     optimizer_type: str = "adam",
     gradient_clip: float | None = None,
+    weight_decay: float = 0.0,
 ) -> optax.GradientTransformation:
-    """Create the benchmark optimizer through the Opifex training substrate."""
+    """Create the benchmark optimizer through the Opifex training substrate.
+
+    ``weight_decay`` is exposed because ``OptimizerConfig`` defaults it to 0.0 while
+    ``optax.adamw`` defaults it to 1e-4; a caller moving off a direct ``optax.adamw``
+    call needs to be able to say which of the two it meant.
+    """
     return create_optimizer(
         OptimizerConfig(
             optimizer_type=optimizer_type,
             learning_rate=learning_rate,
             gradient_clip=gradient_clip,
+            weight_decay=weight_decay,
         )
     )

benchmarks/singlecell/compute_cost.py

@@ -24,9 +24,9 @@
 import jax
 import jax.numpy as jnp
 import numpy as np
-import optax
 from flax import nnx
 
+from benchmarks._optimizers import create_benchmark_optimizer
 from benchmarks.singlecell._gate2_arms import (
     _ProjectionProbe,
     _embedding_probe,
@@ -91,7 +91,14 @@ def measure(
     )
     labels = jax.device_put(rng.integers(0, n_classes, size=effective_batch).astype(np.int32))
     model, forward = _build(arm, n_features, n_components, n_classes, seed)
-    optimizer = nnx.Optimizer(model, optax.adamw(1e-3), wrt=nnx.Param)
+    # Routed through the shared Opifex boundary (benchmarks/_optimizers.py) rather than
+    # optax directly. weight_decay=1e-4 is optax.adamw's own default, restated here
+    # because OptimizerConfig defaults it to 0.0 -- this keeps the optimizer identical.
+    optimizer = nnx.Optimizer(
+        model,
+        create_benchmark_optimizer(learning_rate=1e-3, optimizer_type="adamw", weight_decay=1e-4),
+        wrt=nnx.Param,
+    )
 
     # donate_argnames lets XLA reuse the model/optimizer buffers in place (Flax NNX
     # best practice), so peak memory reflects realistic in-place training, not double

docs/getting-started/installation.md

@@ -1,6 +1,6 @@
 # Installation
 
-DiffBio requires Python 3.11+ and works on Linux, macOS, and Windows (via WSL).
+DiffBio requires Python 3.12+ and works on Linux, macOS, and Windows (via WSL).
 
 ## Quick Install
 

pyproject.toml

@@ -30,7 +30,6 @@ classifiers = [
   "Topic :: Software Development :: Libraries",
   "Topic :: Software Development :: Libraries :: Python Modules",
   "Programming Language :: Python :: 3 :: Only",
-  "Programming Language :: Python :: 3.11",
   "Programming Language :: Python :: 3.12",
   "License :: OSI Approved :: MIT License"
 ]
@@ -46,10 +45,13 @@ dependencies = [
   # JAX ecosystem
   "beartype>=0.14.1",
   "chex>=0.1.7",
-  "flax>=0.12.0",
+  # 0.12.1 is where nnx Variable gained set_value / get_value.
+  "flax>=0.12.1",
   "jax>=0.6.1",
   "jaxtyping>=0.2.20",
-  "optax>=0.1.4",
+  # 0.2.8 stopped setting jax_pmap_shmap_merge, which jax removed in 0.10; below it,
+  # importing flax raises AttributeError and collection dies for the whole suite.
+  "optax>=0.2.8",
   "orbax-checkpoint>=0.11.10",
   # Scientific computing
   "numpy>=1.24",
@@ -65,11 +67,11 @@ keywords = ["jax", "flax", "bioinformatics", "differentiable", "variant-calling"
 license = {file = "LICENSE"}
 name = "diffbio"
 readme = "README.md"
-requires-python = ">=3.11"
+requires-python = ">=3.12"
 version = "0.1.0"
 
 [project.optional-dependencies]
-all = ["diffbio[benchmark,chem,dev,docs,genomics,gpu,soft-ops-advanced,soft-ops-ot,test,torch-io]"]
+all = ["diffbio[benchmark,chem,cuda12,dev,docs,genomics,soft-ops-advanced,soft-ops-ot,test,torch-io]"]
 benchmark = [
   "scib-metrics>=0.5",
   "pynndescent>=0.5",
@@ -79,7 +81,14 @@ benchmark = [
   "pandas>=2.0",
   "matplotlib>=3.7"
 ]
-cuda-dev = ["diffbio[dev,gpu]"]
+# RDKit, for the molecular-graph and fingerprint operators. Every call site imports it
+# lazily, so the rest of diffbio installs and runs without it.
+chem = ["rdkit>=2025.9.3"]
+cuda-dev = ["diffbio[cuda12,dev]"]
+# Named for the CUDA major version, as JAX names its own extras (cuda12, cuda12-local,
+# cuda13); JAX publishes no `gpu` extra. No jaxlib line: every jax cuda extra already
+# requires jaxlib<=X,>=X at its own version, so restating it only lets the two disagree.
+cuda12 = ["jax[cuda12]>=0.6.1"]
 dev = [
   "bandit[toml]>=1.8.6",
   "build>=1.0.3",
@@ -115,11 +124,7 @@ docs = [
   "mkdocstrings-python>=1.1.2",
   "pymdown-extensions>=10.14.3"
 ]
-# RDKit, for the molecular-graph and fingerprint operators. Every call site imports it
-# lazily, so the rest of diffbio installs and runs without it.
-chem = ["rdkit>=2025.9.3"]
 genomics = ["pysam>=0.22.0", "pyfaidx>=0.8.0"]
-gpu = ["jax[cuda12]>=0.6.1", "jaxlib>=0.6.1"]
 soft-ops-advanced = ["optimistix>=0.0.9", "lineax>=0.0.8"]
 soft-ops-ot = ["diffbio[soft-ops-advanced]", "ott-jax>=0.5.0"]
 test = [
@@ -341,7 +346,7 @@ exclude = ["memory-bank", ".deprecated", "examples", ".venv"]
 extend-include = ["*.ipynb"]
 line-length = 100
 src = ["src"]
-target-version = "py311"
+target-version = "py312"
 
 [tool.ruff.format]
 indent-style = "space"
@@ -365,6 +370,15 @@ ignore = [
   "UP015",
   "UP024",
   "UP035",
+  # PEP 695 generic syntax (`def f[T]`, `class C[T]`). Enabled by target-version py312, but
+  # the docstring gate is pydocstyle 6.3.0, whose parser predates PEP 695: it reports D101 /
+  # D103 "missing docstring" for generic functions, and for generic classes with a bound, on
+  # definitions whose docstring is plainly present. 6.3.0 is its final release (2023-01-17),
+  # so no version of it reads this syntax. Re-enable these together with replacing pydocstyle
+  # with ruff's own D rules, which parse it correctly.
+  "UP046", # Generic class uses Generic subclass instead of type parameters
+  "UP047", # Generic function should use type parameters
+  "UP049", # Private type parameter name
   "RUF005"
 ]
 select = ["E", "F", "UP"]
@@ -384,6 +398,11 @@ convention = "google"
 # setuptools arrives through torch, tensorflow and tensorboard, which require it at
 # runtime; it is not this project's build backend, which is hatchling.
 constraint-dependencies = [
+  # dm-haiku arrives through jax-md. Up to 0.0.16 it does
+  # `DropVar: TypeAlias = jax.core.DropVar` at import time, and jax removed jax.core.DropVar
+  # in 0.11.0, so importing diffbio.operators.molecular_dynamics dies with AttributeError.
+  # 0.0.17 guards it with a try/except onto jax.extend.core.
+  "dm-haiku>=0.0.17",
   "setuptools>=83.0.0" # PYSEC-2026-3447
 ]
 python-preference = "only-managed"

setup.sh

@@ -120,7 +120,7 @@ SYNC_ARGS=(sync --extra dev --extra test)
 case "$BACKEND" in
     cpu) ;;
     cuda12)
-        SYNC_ARGS+=(--extra gpu)
+        SYNC_ARGS+=(--extra cuda12)
         ;;
     metal)
         SYNC_ARGS+=(--extra metal)

src/diffbio/core/soft_ops/sorting.py

@@ -990,4 +990,11 @@ def top_k_mask(
     pattern = jnp.concatenate([jnp.ones(k, x.dtype), jnp.zeros(n - k, x.dtype)])
     weights = jnp.broadcast_to(pattern, x_last.shape)
     mask = proj_permutahedron(x_last, weights, softness=softness, mode="c0")
+    # The projection lands in {m : 0 <= m <= 1, sum(m) = k}, so the bound holds in exact
+    # arithmetic. The isotonic solver averages pooled blocks in float32 and can finish one
+    # ULP outside it: for ``n=4`` with scores from seed 202 the largest entry comes back as
+    # 1.0000001. Saturate the forward value onto the documented bound, but leave the backward
+    # pass on the projection's own gradient -- a plain clip would zero the gradient of
+    # precisely the saturated entries, which are the ones the mask selects.
+    mask = mask + jax.lax.stop_gradient(jnp.clip(mask, 0.0, 1.0) - mask)
     return jnp.moveaxis(mask, -1, axis)

src/diffbio/sources/perturbation/_types.py

@@ -8,16 +8,15 @@
 from __future__ import annotations
 
 from enum import StrEnum
-from typing import TypeAlias
 
 # ---------------------------------------------------------------------------
 # Type aliases
 # ---------------------------------------------------------------------------
 
-CellIndex: TypeAlias = int
-PerturbationLabel: TypeAlias = str
-CellTypeLabel: TypeAlias = str
-BatchLabel: TypeAlias = str
+type CellIndex = int
+type PerturbationLabel = str
+type CellTypeLabel = str
+type BatchLabel = str
 
 
 # ---------------------------------------------------------------------------