Remove device parameter, let JAX handle device selection

JAX selects the appropriate device based on its install (CPU/GPU) and
runtime context managers. The explicit device arg added unnecessary
complexity with no benefit over JAX's built-in device management.

Removes device from: align_obs, align_images, AlignBackend protocol,
StAlignBackend, MoscotBackend, and require_jax. Simplifies require_jax
to a pure import guard.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: timtreis

src/squidpy/experimental/tl/_align/_api.py

@@ -48,7 +48,6 @@ def align_obs(
     *,
     output_mode: Literal["affine", "obs", "return"] = "obs",
     key_added: str | None = None,
-    device: Literal["cpu", "gpu"] | None = None,
     inplace: bool = True,
     **flavour_kwargs: Any,
 ) -> AnnData | SpatialData | AlignResult | None:
@@ -83,9 +82,6 @@ def align_obs(
         Name for the aligned table when ``output_mode='obs'`` and inputs are
         SpatialData.  Defaults to ``'{adata_query_name}_aligned'``.
         Rejected with any other ``output_mode``.
-    device
-        ``'cpu'``/``'gpu'`` to force a JAX device, or ``None`` to let JAX
-        pick the default.  Only consulted by JAX-backed flavours.
     inplace
         If ``True``, mutate the query container; otherwise return a copy.
         Only affects SpatialData inputs -- for plain AnnData with
@@ -99,7 +95,7 @@ def align_obs(
 
     pair = resolve_obs_pair(data_ref, data_query, adata_ref_name, adata_query_name)
     backend = get_backend(flavour)
-    result = backend.align_obs(pair, device=device, **flavour_kwargs)
+    result = backend.align_obs(pair, **flavour_kwargs)
 
     # Auto-generate key_added for SpatialData obs writeback.
     if key_added is None and output_mode == "obs" and pair.query_element_key is not None:
@@ -118,7 +114,6 @@ def align_images(
     scale_ref: str | Literal["auto"] = "auto",
     scale_query: str | Literal["auto"] = "auto",
     output_mode: Literal["affine", "return"] = "affine",
-    device: Literal["cpu", "gpu"] | None = None,
     inplace: bool = True,
     **flavour_kwargs: Any,
 ) -> SpatialData | AlignResult | None:
@@ -145,7 +140,7 @@ def align_images(
         ``'affine'`` registers the fit on the query image element so all of
         its scales inherit the transformation; ``'return'`` returns the raw
         :class:`AlignResult`.
-    device, inplace, flavour_kwargs
+    inplace, flavour_kwargs
         See :func:`align_obs`.
     """
     validate_flavour(flavour, allowed=ALLOWED_FLAVOURS_IMAGES, op="align_images")
@@ -160,7 +155,7 @@ def align_images(
         scale_query=scale_query,
     )
     backend = get_backend(flavour)
-    result = backend.align_images(pair, device=device, **flavour_kwargs)
+    result = backend.align_images(pair, **flavour_kwargs)
 
     return _writeback(pair, result, output_mode=output_mode, key_added=None, inplace=inplace)
 

src/squidpy/experimental/tl/_align/_backends/_base.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Literal, Protocol, runtime_checkable
+from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
 
 if TYPE_CHECKING:
     from squidpy.experimental.tl._align._types import AlignPair, AlignResult
@@ -24,15 +24,11 @@ class AlignBackend(Protocol):
     def align_obs(
         self,
         pair: AlignPair,
-        *,
-        device: Literal["cpu", "gpu"] | None = None,
         **kwargs: Any,
     ) -> AlignResult: ...
 
     def align_images(
         self,
         pair: AlignPair,
-        *,
-        device: Literal["cpu", "gpu"] | None = None,
         **kwargs: Any,
     ) -> AlignResult: ...

src/squidpy/experimental/tl/_align/_backends/_moscot.py

@@ -8,7 +8,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Literal
+from typing import TYPE_CHECKING, Any
 
 if TYPE_CHECKING:
     from squidpy.experimental.tl._align._types import AlignPair, AlignResult
@@ -21,13 +21,11 @@ class MoscotBackend:
     def align_obs(
         self,
         pair: AlignPair,
-        *,
-        device: Literal["cpu", "gpu"] | None = None,
         **kwargs: Any,
     ) -> AlignResult:
         from squidpy.experimental.tl._align._jax import require_jax
 
-        require_jax(device)
+        require_jax()
         raise NotImplementedError(
             "moscot backend `align_obs`: TODO. Skeleton landed; the moscot "
             "solver will replace this body in a follow-up PR."
@@ -36,8 +34,6 @@ def align_obs(
     def align_images(
         self,
         pair: AlignPair,
-        *,
-        device: Literal["cpu", "gpu"] | None = None,
         **kwargs: Any,
     ) -> AlignResult:
         raise NotImplementedError("moscot does not implement image alignment; use `flavour='stalign'`.")

src/squidpy/experimental/tl/_align/_backends/_stalign.py

@@ -1,13 +1,13 @@
 """STalign backend.
 
-Wraps the JAX LDDMM solver lifted from scverse/squidpy#1150 (Selman Özleyen)
+Wraps the JAX LDDMM solver lifted from scverse/squidpy#1150 (Selman Ozleyen)
 into the :class:`AlignBackend` Protocol. Only ``align_obs`` is implemented
 today; ``align_images`` raises until upstream support exists.
 """
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Literal
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 
@@ -23,7 +23,6 @@ def align_obs(
         self,
         pair: AlignPair,
         *,
-        device: Literal["cpu", "gpu"] | None = None,
         config: Any | None = None,
         landmarks_source: np.ndarray | None = None,
         landmarks_target: np.ndarray | None = None,
@@ -39,7 +38,7 @@ def align_obs(
         # `ModuleNotFoundError: import of jax halted; None in sys.modules`
         # instead of the clean `ImportError("JAX is required ...")` from
         # _jax.require_jax.
-        require_jax(device)
+        require_jax()
 
         from squidpy.experimental.tl._align._backends._stalign_tools import stalign_points
         from squidpy.experimental.tl._align._types import AlignResult, ObsDisplacement
@@ -94,8 +93,6 @@ def align_obs(
     def align_images(
         self,
         pair: AlignPair,
-        *,
-        device: Literal["cpu", "gpu"] | None = None,
         **kwargs: Any,
     ) -> AlignResult:
         raise NotImplementedError(

src/squidpy/experimental/tl/_align/_jax.py

@@ -1,4 +1,4 @@
-"""Lazy JAX import + device selection for JAX-backed alignment backends.
+"""Lazy JAX import guard for JAX-backed alignment backends.
 
 JAX is an optional dependency.  Importing this module is cheap; calling
 :func:`require_jax` is what actually pulls JAX in, and only the
@@ -7,50 +7,22 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Literal
+from typing import Any
 
-if TYPE_CHECKING:
-    Device = Any  # jax.Device, but importing it eagerly defeats the purpose
+_INSTALL_HINT = 'JAX is required for the requested align_* flavour. Install with `pip install "squidpy[jax]"`.'
 
 
-_INSTALL_HINT = (
-    "JAX is required for the requested align_* flavour. "
-    "Install with `pip install jax` (CPU) or follow the JAX install guide for GPU."
-)
-
-
-def require_jax(device: Literal["cpu", "gpu"] | None = None) -> tuple[Any, Any]:
-    """Import JAX lazily and return ``(jax, device)``.
-
-    Parameters
-    ----------
-    device
-        ``"cpu"``/``"gpu"`` to force a platform, or ``None`` to use whatever
-        JAX picks as the default.
-
-    Returns
-    -------
-    jax_module
-        The imported :mod:`jax` module.
-    device
-        A :class:`jax.Device` of the requested platform.
+def require_jax() -> Any:
+    """Import JAX lazily and return the module.
 
     Raises
     ------
     ImportError
         If JAX is not installed.
-    RuntimeError
-        If the requested device platform is not available on this host.
     """
     try:
         import jax
     except ImportError as e:
         raise ImportError(_INSTALL_HINT) from e
 
-    if device is None:
-        return jax, jax.devices()[0]
-
-    matching = [d for d in jax.devices() if d.platform == device]
-    if not matching:
-        raise RuntimeError(f"No JAX device of kind {device!r} available; have {[d.platform for d in jax.devices()]}.")
-    return jax, matching[0]
+    return jax