diff --git a/src/quantem/core/config.py b/src/quantem/core/config.py
index c2eda10ba..7213f3b5f 100644
--- a/src/quantem/core/config.py
+++ b/src/quantem/core/config.py
@@ -34,6 +34,15 @@
     if "cuda" in str(e):
         NUM_DEVICES = 0
     _defaults["has_cupy"] = False
+try:
+    import quantem.cuda  # type: ignore  # noqa: F401
+
+    _defaults["has_quantem_cuda"] = True
+except ModuleNotFoundError:
+    _defaults["has_quantem_cuda"] = False
+except Exception:
+    # installed but unloadable (e.g. libcudart missing at runtime)
+    _defaults["has_quantem_cuda"] = False
 
 
 defaults: list[Mapping] = [_defaults]
diff --git a/src/quantem/core/ml/activation_functions.py b/src/quantem/core/ml/activation_functions.py
index 25db3edbe..1f3675660 100644
--- a/src/quantem/core/ml/activation_functions.py
+++ b/src/quantem/core/ml/activation_functions.py
@@ -1,10 +1,66 @@
-from typing import Callable
+import os
+from typing import Any, Callable
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
 
+class _TruncExpReference(torch.autograd.Function):
+    """Exponential with the torch-ngp clamped backward used by tomography."""
+
+    @staticmethod
+    def forward(ctx: Any, values: torch.Tensor, offset: float) -> torch.Tensor:
+        ctx.save_for_backward(values)
+        ctx.offset = offset
+        return torch.exp(values - offset)
+
+    @staticmethod
+    def backward(ctx: Any, grad_output: torch.Tensor) -> tuple[torch.Tensor, None]:
+        (values,) = ctx.saved_tensors
+        shifted = values - ctx.offset
+        return grad_output * torch.exp(shifted.clamp(max=15)), None
+
+
+def trunc_exp(values: torch.Tensor, offset: float = 0.0) -> torch.Tensor:
+    """Apply trunc-exp, using the fused CUDA density tail when available."""
+    use_fused = (
+        os.environ.get("QUANTEM_DENSITY_TAIL_FUSED", "1") != "0"
+        and values.is_cuda
+        and values.dtype in (torch.float32, torch.bfloat16)
+        and values.ndim in (1, 2)
+        and values.numel() > 0
+    )
+    if use_fused:
+        try:
+            import quantem.cuda.core.ml as cuda_ml
+        except (ImportError, OSError, RuntimeError):
+            pass
+        else:
+            fused = getattr(cuda_ml, "density_tail", None)
+            if fused is not None:
+                return fused(values, float(offset))
+    return _TruncExpReference.apply(values, float(offset))
+
+
+# Object constraints consume this explicit capability instead of guessing from
+# a callable name or implementation detail.
+trunc_exp.quantem_guarantees_nonnegative = True  # type: ignore[attr-defined]
+
+
+class TruncExpActivation(nn.Module):
+    """Configurable trunc-exp activation with a non-negativity capability."""
+
+    quantem_guarantees_nonnegative = True
+
+    def __init__(self, offset: float = 0.0) -> None:
+        super().__init__()
+        self.offset = float(offset)
+
+    def forward(self, values: torch.Tensor) -> torch.Tensor:
+        return trunc_exp(values, self.offset)
+
+
 class ModReLU(nn.Module):
     """Modulated ReLU activation for complex-valued inputs.
 
diff --git a/src/quantem/core/ml/inr.py b/src/quantem/core/ml/inr.py
index 4c1973447..f48017057 100644
--- a/src/quantem/core/ml/inr.py
+++ b/src/quantem/core/ml/inr.py
@@ -111,18 +111,20 @@ def _build(self) -> None:
         self.net = nn.Sequential(*net_list)
 
         if self.winner_initialization:
-            if type(self.winner_initialization) is int:
-                rng = torch.Generator()
-                rng.manual_seed(self.winner_initialization)
-            else:
-                rng = torch.Generator()
-                rng.manual_seed(42)
+            seed = self.winner_initialization if type(self.winner_initialization) is int else 42
+            rng = torch.Generator()
+            rng.manual_seed(seed)
+            # torch.randn_like ignores generators, so the noise must come from
+            # torch.randn with the seeded generator -- otherwise the "winner" seed
+            # silently has no effect and the perturbation is not reproducible.
             with torch.no_grad():
-                self.net[0].linear.weight += (  # type: ignore[reportAttributeAccessIssue]
-                    torch.randn_like(self.net[0].linear.weight) * 5 / self.first_omega_0  # type:ignore
-                )
-                self.net[1].linear.weight += (  # type: ignore[reportAttributeAccessIssue]
-                    torch.randn_like(self.net[1].linear.weight) * 0.1 / self.hidden_omega_0  # type:ignore
+                w0 = self.net[0].linear.weight  # type: ignore[reportAttributeAccessIssue]
+                w0 += torch.randn(w0.shape, generator=rng, dtype=w0.dtype) * 5 / self.first_omega_0
+                w1 = self.net[1].linear.weight  # type: ignore[reportAttributeAccessIssue]
+                w1 += (
+                    torch.randn(w1.shape, generator=rng, dtype=w1.dtype)
+                    * 0.1
+                    / self.hidden_omega_0
                 )
 
     def forward(self, coords: torch.Tensor) -> torch.Tensor:
diff --git a/src/quantem/core/ml/logger.py b/src/quantem/core/ml/logger.py
index 13f3beb73..ca5875485 100644
--- a/src/quantem/core/ml/logger.py
+++ b/src/quantem/core/ml/logger.py
@@ -5,7 +5,7 @@
 import shutil
 import tempfile
 from pathlib import Path
-from typing import Self, cast
+from typing import Any, Literal, Mapping, Self, cast
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -16,11 +16,21 @@
 
 from quantem.core.io.serialize import AutoSerialize, load
 
-"""Tensorboard logger class for AD/ML reconstruction methods."""
+"""Logger class for AD/ML reconstruction methods."""
 
 
 class LoggerBase(AutoSerialize):
-    """Tensorboard logger for AD/ML reconstruction methods."""
+    """Logger for AD/ML reconstruction methods.
+
+    Parameters
+    ----------
+    mode : {"tensorboard", "wandb"}, optional
+        Logging backend. TensorBoard preserves the historical behavior. WandB runs default to
+        offline mode unless ``WANDB_MODE`` is already set, and their files are written under
+        ``<log_dir>/wandb/``. Upload offline runs later with ``wandb sync <log_dir>/wandb/<run>``.
+    wandb_config : Mapping, optional
+        Configuration attached to ``wandb.init(config=...)``. Ignored by TensorBoard mode.
+    """
 
     def __init__(
         self,
@@ -28,6 +38,8 @@ def __init__(
         run_prefix: str,
         run_suffix: str = "",
         log_images_every: int = 10,
+        mode: Literal["tensorboard", "wandb"] | str = "tensorboard",
+        wandb_config: Mapping[str, Any] | None = None,
     ) -> None:
         """Initialize LoggerBase.
 
@@ -41,27 +53,93 @@ def __init__(
             Suffix for run directory name, by default ""
         log_images_every : int, optional
             Frequency for logging images, by default 10
+        mode : {"tensorboard", "wandb"}, optional
+            Logging backend, by default "tensorboard"
+        wandb_config : Mapping, optional
+            Configuration attached to WandB runs. No-op for TensorBoard.
         """
         self._timestamp = datetime.datetime.now().strftime(
             "%Y%m%d_%H%M%S"
         )  # This should never be reinstantiated.
+        # DDP: only the main rank owns the writer and the on-disk run directory.
+        # Non-main ranks build a fully inert logger -- no SummaryWriter/WandB run, no
+        # directory created, and every log_*/flush/close call is a no-op -- so a DDP
+        # run yields a single run subdirectory instead of one per rank. RANK is unset
+        # for single-process runs, so this defaults to rank 0 and preserves behavior.
+        self._is_writer_rank = int(os.environ.get("RANK", "0")) == 0
         self.run_prefix = run_prefix
         self.run_suffix = run_suffix
         self.log_dir = base_log_dir
         self.log_images_every = log_images_every
-        self.writer = SummaryWriter(str(self.log_dir))
-
-    def log_scalar(self, tag: str, value: float, step: int) -> None:
-        self.writer.add_scalar(tag=tag, scalar_value=value, global_step=step)
-
-    def log_image(self, tag: str, image: NDArray | Tensor, step: int, cmap: str = "turbo") -> None:
+        self.mode = mode
+        self._wandb_run = None
+        self._wandb = None
+        self.writer = None
+
+        if self._is_writer_rank:
+            if self.mode == "tensorboard":
+                self.writer = SummaryWriter(str(self.log_dir))
+            else:
+                self._init_wandb(wandb_config)
+
+    def log_scalar(
+        self,
+        tag: str,
+        value: float,
+        step: int,
+        step_domain: str = "epoch",
+        extra_steps: dict[str, int] | None = None,
+    ) -> None:
+        if not self._is_writer_rank:
+            return
+        if self.mode == "tensorboard":
+            self.writer.add_scalar(tag=tag, scalar_value=value, global_step=step)
+        else:
+            self._log_wandb(tag, float(value), step, step_domain, extra_steps=extra_steps)
+
+    def log_image(
+        self,
+        tag: str,
+        image: NDArray | Tensor,
+        step: int,
+        cmap: str = "turbo",
+        step_domain: str = "epoch",
+        extra_steps: dict[str, int] | None = None,
+    ) -> None:
+        if not self._is_writer_rank:
+            return
         cmap_image = self.apply_colormap(image, cmap_name=cmap)
-        self.writer.add_image(tag, cmap_image, step)
-
-    def log_figure(self, tag: str, fig: Figure, step: int) -> None:
-        self.writer.add_figure(tag, fig, step)
-
-    def log_histogram(self, tag: str, values: NDArray | Tensor, step: int) -> None:
+        if self.mode == "tensorboard":
+            self.writer.add_image(tag, cmap_image, step)
+        else:
+            image_hwc = np.moveaxis(cmap_image, 0, -1)
+            self._log_wandb(
+                tag, self._wandb.Image(image_hwc), step, step_domain, extra_steps=extra_steps
+            )
+
+    def log_figure(
+        self,
+        tag: str,
+        fig: Figure,
+        step: int,
+        step_domain: str = "epoch",
+        extra_steps: dict[str, int] | None = None,
+    ) -> None:
+        if not self._is_writer_rank:
+            return
+        if self.mode == "tensorboard":
+            self.writer.add_figure(tag, fig, step)
+        else:
+            self._log_wandb(tag, self._wandb.Image(fig), step, step_domain, extra_steps=extra_steps)
+
+    def log_histogram(
+        self,
+        tag: str,
+        values: NDArray | Tensor,
+        step: int,
+        step_domain: str = "epoch",
+        extra_steps: dict[str, int] | None = None,
+    ) -> None:
         """Log histogram of values for monitoring distributions.
 
         Parameters
@@ -73,11 +151,25 @@ def log_histogram(self, tag: str, values: NDArray | Tensor, step: int) -> None:
         step : int
             Step number.
         """
+        if not self._is_writer_rank:
+            return
         if isinstance(values, Tensor):
             values = values.detach().cpu().numpy()
-        self.writer.add_histogram(tag, values, step)
-
-    def log_text(self, tag: str, text: str, step: int) -> None:
+        if self.mode == "tensorboard":
+            self.writer.add_histogram(tag, values, step)
+        else:
+            self._log_wandb(
+                tag, self._wandb.Histogram(values), step, step_domain, extra_steps=extra_steps
+            )
+
+    def log_text(
+        self,
+        tag: str,
+        text: str,
+        step: int,
+        step_domain: str = "epoch",
+        extra_steps: dict[str, int] | None = None,
+    ) -> None:
         """Log text for configuration, hyperparameters, or notes.
 
         Parameters
@@ -89,14 +181,36 @@ def log_text(self, tag: str, text: str, step: int) -> None:
         step : int
             Step number.
         """
-        self.writer.add_text(tag, text, step)
+        if not self._is_writer_rank:
+            return
+        if self.mode == "tensorboard":
+            self.writer.add_text(tag, text, step)
+        else:
+            self._log_wandb(tag, text, step, step_domain, extra_steps=extra_steps)
+
+    def attach_config(self, config: Mapping[str, Any]) -> None:
+        """Attach a resolved run configuration to WandB.
+
+        This is intentionally a no-op for TensorBoard mode so callers can use it unconditionally.
+        """
+        if not self._is_writer_rank:
+            return
+        if self.mode == "wandb":
+            self._wandb_run.config.update(dict(config), allow_val_change=True)
 
     def flush(self) -> None:
-        self.writer.flush()
+        if self._is_writer_rank and self.mode == "tensorboard":
+            self.writer.flush()
 
     def close(self) -> None:
-        self.writer.flush()
-        self.writer.close()
+        if not self._is_writer_rank:
+            return
+        if self.mode == "tensorboard":
+            self.writer.flush()
+            self.writer.close()
+        elif self._wandb_run is not None:
+            self._wandb_run.finish()
+            self._wandb_run = None
 
     def new_timestamp(self) -> None:
         """Create new timestamp and reinitialize writer with new log directory."""
@@ -106,9 +220,14 @@ def new_timestamp(self) -> None:
         if self.run_suffix:
             name += f"_{self.run_suffix}"
         new_log_dir = self.log_dir.parent / name
-        new_log_dir.mkdir(exist_ok=True)
         self._log_dir = new_log_dir
-        self.writer = SummaryWriter(str(self.log_dir))
+        if not self._is_writer_rank:
+            return
+        new_log_dir.mkdir(exist_ok=True)
+        if self.mode == "tensorboard":
+            self.writer = SummaryWriter(str(self.log_dir))
+        else:
+            self._init_wandb()
 
     def clone(self) -> Self:
         """Create a cloned logger with new timestamp.
@@ -158,7 +277,10 @@ def log_dir(self, dir: str | os.PathLike) -> None:
             name += f"_{self.run_suffix}"
 
         full_path = dir / name
-        full_path.mkdir(parents=True, exist_ok=True)
+        # Non-main DDP ranks never materialize a run directory (see __init__). getattr
+        # guards the first setter call, which runs before _is_writer_rank is assigned.
+        if getattr(self, "_is_writer_rank", True):
+            full_path.mkdir(parents=True, exist_ok=True)
 
         self._log_dir = full_path
 
@@ -192,6 +314,16 @@ def log_images_every(self) -> int:
     def log_images_every(self, value: int) -> None:
         self._log_images_every = int(value)
 
+    @property
+    def mode(self) -> str:
+        return self._mode
+
+    @mode.setter
+    def mode(self, value: str) -> None:
+        if value not in {"tensorboard", "wandb"}:
+            raise ValueError("Logger mode must be 'tensorboard' or 'wandb'.")
+        self._mode = value
+
     # --- Helper Functions ---
 
     @staticmethod
@@ -217,3 +349,53 @@ def apply_colormap(tensor_2d: Tensor | NDArray, cmap_name: str = "turbo") -> NDA
         cmap = plt.get_cmap(cmap_name)
         colored = cmap(tensor_2d)[..., :3].transpose(2, 0, 1)  # type: ignore # [3, H, W]
         return colored.astype(np.float32)
+
+    def _init_wandb(self, config: Mapping[str, Any] | None = None) -> None:
+        """Initialize a WandB run rooted inside this logger's run directory."""
+        os.environ.setdefault("WANDB_MODE", "offline")
+        wandb_dir = self.log_dir / "wandb"
+        wandb_dir.mkdir(parents=True, exist_ok=True)
+        for env_name, path in {
+            "WANDB_DIR": wandb_dir,
+            "WANDB_CONFIG_DIR": wandb_dir / "config",
+            "WANDB_CACHE_DIR": wandb_dir / "cache",
+            "WANDB_DATA_DIR": wandb_dir / "data",
+        }.items():
+            os.environ[env_name] = str(path)
+            path.mkdir(parents=True, exist_ok=True)
+
+        try:
+            import wandb
+        except ImportError as exc:
+            raise ImportError(
+                "LoggerBase(mode='wandb') requires the 'wandb' package. "
+                "The NERSC quantem environment at "
+                "/global/common/software/m5020/cedlim/conda/quantem/bin/python3 "
+                "includes wandb 0.28.0."
+            ) from exc
+
+        self._wandb = wandb
+        self._wandb_run = wandb.init(
+            project="quantem-tomography",
+            name=self.log_dir.name,
+            dir=str(self.log_dir),
+            config=dict(config) if config is not None else None,
+            reinit=True,
+        )
+        self._wandb_run.define_metric("*", step_metric="epoch")
+        self._wandb_run.define_metric("snapshots/*", step_metric="grad_step")
+
+    def _log_wandb(
+        self,
+        tag: str,
+        value: Any,
+        step: int,
+        step_domain: str,
+        extra_steps: dict[str, int] | None = None,
+    ) -> None:
+        """Log WandB data with named step metrics instead of the global step."""
+        payload = {tag: value}
+        if extra_steps is not None:
+            payload.update(extra_steps)
+        payload[step_domain] = step
+        self._wandb_run.log(payload)
diff --git a/src/quantem/core/ml/models/kplanes.py b/src/quantem/core/ml/models/kplanes.py
index cdf552617..26991521c 100644
--- a/src/quantem/core/ml/models/kplanes.py
+++ b/src/quantem/core/ml/models/kplanes.py
@@ -3,6 +3,9 @@
 """
 
 import itertools
+import os
+import threading
+import warnings
 from typing import Callable, Optional, Sequence
 
 # import tinycudann as tcnn
@@ -10,9 +13,112 @@
 import torch.nn.functional as F
 from torch import nn
 
+from quantem.core import config
+
 from .model_base import PPLR, TensorDecompositionModel
 from .so3params import SO3ParamQuat, SO3ParamR9SVD
 
+_FusedMlpShapeKey = tuple[int | None, int, int, int, int, int]
+_unsupported_fused_mlp_shapes: dict[_FusedMlpShapeKey, str] = {}
+_warned_fused_mlp_reasons: set[str] = set()
+_fused_mlp_fallback_lock = threading.Lock()
+
+
+def _fused_mlp_enabled() -> bool:
+    return os.environ.get("QUANTEM_FUSED_MLP", "1") != "0"
+
+
+class FusedHiddenMLP(nn.Sequential):
+    """Sequential-compatible sigma head with a default-on cuBLASLt fast path.
+
+    Keeping the numeric child-module names preserves existing state dicts and
+    optimizer parameter discovery.  Every unsupported case executes the
+    inherited Sequential path without importing quantem-cuda eagerly.
+    ``torch.use_deterministic_algorithms`` does not govern the custom path;
+    set ``QUANTEM_FUSED_MLP=0`` when strict PyTorch deterministic-mode behavior
+    is required.
+    """
+
+    def _can_fuse(self, inputs: torch.Tensor) -> bool:
+        if not _fused_mlp_enabled():
+            return False
+        if torch.compiler.is_compiling():
+            return False
+        if (
+            not inputs.is_cuda
+            or inputs.dtype != torch.bfloat16
+            or not inputs.is_contiguous()
+            or not torch.is_autocast_enabled("cuda")
+            or torch.get_autocast_dtype("cuda") != torch.bfloat16
+        ):
+            return False
+        if len(self) != 5:
+            return False
+        linears = (self[0], self[2], self[4])
+        if not (
+            all(isinstance(layer, nn.Linear) for layer in linears)
+            and isinstance(self[1], nn.ReLU)
+            and isinstance(self[3], nn.ReLU)
+            and all(layer.bias is not None for layer in linears)
+        ):
+            return False
+        parameters = tuple(
+            parameter for layer in linears for parameter in (layer.weight, layer.bias)
+        )
+        return all(
+            parameter.device == inputs.device
+            and parameter.dtype == torch.float32
+            and parameter.is_contiguous()
+            for parameter in parameters
+        )
+
+    def _fused_shape_key(self, inputs: torch.Tensor) -> _FusedMlpShapeKey:
+        linear1, linear2, linear3 = self[0], self[2], self[4]
+        return (
+            inputs.device.index,
+            inputs.shape[0],
+            inputs.shape[1],
+            linear1.out_features,
+            linear2.out_features,
+            linear3.out_features,
+        )
+
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        if self._can_fuse(inputs):
+            shape_key = self._fused_shape_key(inputs)
+            with _fused_mlp_fallback_lock:
+                shape_is_unsupported = shape_key in _unsupported_fused_mlp_shapes
+            if shape_is_unsupported:
+                return super().forward(inputs)
+            try:
+                from quantem.cuda.core.ml import fused_hidden_mlp
+
+                linear1, linear2, linear3 = self[0], self[2], self[4]
+                return fused_hidden_mlp(
+                    inputs,
+                    linear1.weight,
+                    linear1.bias,
+                    linear2.weight,
+                    linear2.bias,
+                    linear3.weight,
+                    linear3.bias,
+                )
+            except (ImportError, AttributeError, RuntimeError, TypeError, ValueError) as error:
+                reason = f"{type(error).__name__}: {error}"
+                with _fused_mlp_fallback_lock:
+                    _unsupported_fused_mlp_shapes[shape_key] = reason
+                    should_warn = reason not in _warned_fused_mlp_reasons
+                    _warned_fused_mlp_reasons.add(reason)
+                if should_warn:
+                    warnings.warn(
+                        "Fused cuBLASLt MLP dispatch failed; memoizing the eager fallback for "
+                        f"this device/shape. Reason: {reason}",
+                        RuntimeWarning,
+                        stacklevel=2,
+                    )
+        return super().forward(inputs)
+
+
 """
 K-planes utility functions
 """
@@ -99,6 +205,7 @@ def init_planes(
             nn.init.ones_(param)
         else:
             nn.init.uniform_(param, a=a, b=b)
+        param.data = param.data.contiguous(memory_format=torch.channels_last)
         planes.append(param)
     return planes
 
@@ -144,12 +251,15 @@ def interpolate_ms_features(
     pts: torch.Tensor,
     ms_grids: nn.ParameterList,
 ) -> torch.Tensor:
-    mat_mode = [[0, 1], [0, 2], [1, 2]]
+    # Plane axis layout: XY=(0,1), XZ=(0,2), YZ=(1,2). Stacking views avoids the
+    # per-call index-tensor allocation (a host-to-device copy) that list-based
+    # advanced indexing does three times per forward.
+    x, y, z = pts.unbind(-1)
     coord_plane = torch.stack(
         [
-            pts[:, mat_mode[0]],
-            pts[:, mat_mode[1]],
-            pts[:, mat_mode[2]],
+            torch.stack((x, y), dim=-1),
+            torch.stack((x, z), dim=-1),
+            torch.stack((y, z), dim=-1),
         ]
     ).view(3, -1, 1, 2)
 
@@ -169,6 +279,7 @@ class KPlanes(PPLR, TensorDecompositionModel):
     """
     K-Planes model adapted from Fridovich-Keil et al., https://arxiv.org/abs/2301.10241
     """
+
     def __init__(
         self,
         # Grid parameters
@@ -199,12 +310,22 @@ def __init__(
         self.concat_features = concat_features
         self.density_activation = density_activation
 
+        # All three planes share one (3, C, res[1], res[0]) tensor, which ignores
+        # res[2]: an anisotropic resolution would silently give the XZ/YZ planes the
+        # wrong grid along z. Refuse it rather than misallocate.
+        if len(set(self.resolution)) != 1:
+            raise ValueError(
+                f"KPlanes currently requires an isotropic resolution, got {self.resolution}; "
+                "the plane grids are allocated as (res[1], res[0]) for all three axis pairs."
+            )
+
         self.grids = nn.ParameterList()
         self.feature_dim = 0
         for res_mult in self.multiscale_res_multipliers:
             scaled_res = [int(r * res_mult) for r in self.resolution]
             plane = nn.Parameter(torch.empty(3, self.M_features, scaled_res[1], scaled_res[0]))
             nn.init.uniform_(plane, 0.1, 0.5)
+            plane.data = plane.data.contiguous(memory_format=torch.channels_last)
             self.grids.append(plane)
             self.feature_dim += self.M_features
 
@@ -232,7 +353,20 @@ def __init__(
             nn.init.normal_(out.weight, std=0.01)
             nn.init.zeros_(out.bias)
             layers.append(out)
-            self.sigma_net = nn.Sequential(*layers)
+            self.sigma_net = FusedHiddenMLP(*layers)
+        else:
+            # Linear head fallback, matching KPlanesTILTED._build_sigma_net and
+            # CPTilted: forward/get_params reference sigma_net unconditionally.
+            self.sigma_net = nn.Linear(self.feature_dim, 1, bias=True)
+            nn.init.normal_(self.sigma_net.weight, std=0.01)
+            nn.init.zeros_(self.sigma_net.bias)
+
+    def __setstate__(self, state) -> None:
+        """Restore channels-last grids when loading legacy whole-module checkpoints."""
+        super().__setstate__(state)
+        for plane in self.grids:
+            if plane.ndim == 4:
+                plane.data = plane.data.contiguous(memory_format=torch.channels_last)
 
     def get_densities(self, coords: torch.Tensor):
         """Computes and returns densities"""
@@ -278,16 +412,6 @@ def tilted(self, tilted: bool):
             raise TypeError("tilted must be a boolean")
         self._tilted = tilted
 
-    @property
-    def grids(self) -> torch.nn.ParameterList:
-        return self._grids
-
-    @grids.setter
-    def grids(self, grids: torch.nn.ParameterList):
-        if not isinstance(grids, torch.nn.ParameterList):
-            raise TypeError("Grids must be a ParameterList")
-        self._grids = grids
-
     @property
     def resolution(self) -> list[int]:
         return self._resolution
@@ -308,32 +432,102 @@ def interpolate_ms_features_tilted(
     pts: torch.Tensor,  # (B, 3)
     ms_grids: nn.ParameterList,  # each grid: (3*T, C, H, W)
     rotation_matrices: torch.Tensor,  # (T, 3, 3)
-) -> torch.Tensor:
+    scale_gates: Optional[Sequence[float]] = None,  # per-scale multiplier (coarse-to-fine)
+    include_plane_tv: bool = False,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
     """
     Fully-vectorized multi-scale, multi-rotation K-Planes feature interpolation.
     Returns features of shape (B, C * T * num_scales).
+
+    scale_gates (optional): one scalar per scale in [0,1] applied to that scale's feature
+    block before concatenation. Used for coarse-to-fine band-limiting: ramp the finer
+    scales on over training so the model commits to view-consistent low-frequency structure
+    first (a classical limited-angle prior). None => all scales fully on (default behaviour).
+
+    When the optional ``quantem-cuda`` package is installed, the tensors live
+    on a CUDA device, and the ``use_cuda_kernels`` config option is true
+    (default), a three-scale model dispatches to one multiscale CUDA op
+    (rotate + grid_sample + Hadamard product + scale gate, analytic backward).
+    Older ``quantem-cuda`` packages and other scale counts fall back to the
+    per-level fused op. Both paths preserve the torch semantics below.
     """
     T = rotation_matrices.shape[0]
     B = pts.shape[0]
 
+    if (
+        pts.is_cuda
+        and pts.dtype == torch.float32
+        and rotation_matrices.dtype == torch.float32
+        and all(g.dtype == torch.float32 for g in ms_grids)
+        and config.get("has_quantem_cuda")
+        and config.get("use_cuda_kernels", default=True)
+    ):
+        import quantem.cuda.core.ml as cuda_ml
+
+        kplanes_tilted_fuse = cuda_ml.kplanes_tilted_fuse
+        kplanes_tilted_fuse_ms = getattr(cuda_ml, "kplanes_tilted_fuse_ms", None)
+        kplanes_tilted_fuse_ms_tv = getattr(cuda_ml, "kplanes_tilted_fuse_ms_tv", None)
+        # Respect instrumentation/overrides of the public single-level op.
+        # The built-in identity is absent in older packages, which naturally
+        # selects the compatible per-level path too.
+        single_level_is_builtin = (
+            getattr(cuda_ml, "_kplanes_tilted_fuse_builtin", None) is kplanes_tilted_fuse
+        )
+        if kplanes_tilted_fuse_ms is not None and len(ms_grids) == 3 and single_level_is_builtin:
+            gates = scale_gates if scale_gates is not None else (1.0, 1.0, 1.0)
+            if (
+                include_plane_tv
+                and os.environ.get("QUANTEM_KPLANES_MS_TV_FUSED", "1") != "0"
+                and kplanes_tilted_fuse_ms_tv is not None
+            ):
+                return kplanes_tilted_fuse_ms_tv(
+                    pts,
+                    rotation_matrices,
+                    ms_grids[0],
+                    ms_grids[1],
+                    ms_grids[2],
+                    float(gates[0]),
+                    float(gates[1]),
+                    float(gates[2]),
+                )
+            return kplanes_tilted_fuse_ms(
+                pts,
+                rotation_matrices,
+                ms_grids[0],
+                ms_grids[1],
+                ms_grids[2],
+                float(gates[0]),
+                float(gates[1]),
+                float(gates[2]),
+            )
+
+        feats = [kplanes_tilted_fuse(pts, rotation_matrices, g) for g in ms_grids]
+        if scale_gates is not None:
+            feats = [f * scale_gates[si] for si, f in enumerate(feats)]
+        return torch.cat(feats, dim=-1)
+
     # (T, B, 3)  — rotate all points by all rotations at once
     rotated = torch.einsum("tij,bj->tbi", rotation_matrices, pts)
 
     # Build (T, 3, B, 2) coords for planes XY, ZX, YZ in one shot.
-    # index_select is faster and cleaner than advanced indexing with python lists.
-    # Plane axis layout: XY=(0,1), ZX=(2,0), YZ=(1,2)
-    idx = torch.tensor([[0, 1], [2, 0], [1, 2]], device=pts.device)  # (3, 2)
-    # rotated: (T, B, 3) -> gather along last dim with idx (3, 2)
-    # Result: (T, 3, B, 2)
-    coords = (
-        rotated.unsqueeze(1).expand(T, 3, B, 3).gather(-1, idx.view(1, 3, 1, 2).expand(T, 3, B, 2))
-    )
+    # Plane axis layout: XY=(0,1), ZX=(2,0), YZ=(1,2). Stacking views avoids the
+    # per-call index-tensor allocation (a host-to-device copy) and the
+    # (T, 3, B, 3) expand+gather this used to do.
+    x, y, z = rotated.unbind(-1)  # each (T, B)
+    coords = torch.stack(
+        (
+            torch.stack((x, y), dim=-1),
+            torch.stack((z, x), dim=-1),
+            torch.stack((y, z), dim=-1),
+        ),
+        dim=1,
+    )  # (T, 3, B, 2)
 
     # Flatten (T, 3) -> 3*T so it matches grid's first dim, and add the H_out=1 axis
     coord_tensor = coords.reshape(3 * T, B, 1, 2)  # (3T, B, 1, 2)
 
     per_scale_features = []
-    for plane_coef in ms_grids:
+    for si, plane_coef in enumerate(ms_grids):
         # plane_coef: (3T, C, H, W)
         C = plane_coef.shape[1]
 
@@ -349,7 +543,10 @@ def interpolate_ms_features_tilted(
         sampled = sampled.squeeze(-1).view(T, 3, C, B).prod(dim=1)
 
         # (T, C, B) -> (B, T, C) -> (B, T*C) to concatenate rotations along feature dim
-        per_scale_features.append(sampled.permute(2, 0, 1).reshape(B, T * C))
+        feat = sampled.permute(2, 0, 1).reshape(B, T * C)
+        if scale_gates is not None:
+            feat = feat * scale_gates[si]
+        per_scale_features.append(feat)
 
     # Concatenate across scales -> (B, T * C * num_scales)
     return torch.cat(per_scale_features, dim=-1)
@@ -389,6 +586,8 @@ class KPlanesTILTED(KPlanes):
     All other args are forwarded to KPlanes.
     """
 
+    quantem_supports_fused_plane_tv = True
+
     def __init__(
         self,
         # Grid parameters
@@ -405,6 +604,9 @@ def __init__(
         hybrid_hidden_dim: int = 64,
         hybrid_num_layers: int = 2,
         so3_param_type: str = "r9svd",
+        # Coarse-to-fine band-limiting: if > 0, finer scales are gated on over the first
+        # c2f_warmup_frac of training via set_progress() (0 => disabled, all scales on).
+        c2f_warmup_frac: float = 0.0,
     ):
         self._td_type = "tilted"
         if input_coords_dims != 3:
@@ -442,6 +644,7 @@ def __init__(
             scaled_res = [int(r * res_mult) for r in resolution]
             plane = nn.Parameter(torch.empty(3 * T, M_features, scaled_res[1], scaled_res[0]))
             nn.init.uniform_(plane, 0.1, 0.5)
+            plane.data = plane.data.contiguous(memory_format=torch.channels_last)
             self.grids.append(plane)
 
         # ---- Rebuild sigma_net with the correct feature_dim ----
@@ -452,6 +655,26 @@ def __init__(
 
         # ---- Learnable rotations ----
         self.set_so3_param_type(so3_param_type, init=tau_init)
+        self.register_buffer("_rotation_matrices_override", None, persistent=False)
+
+        # ---- Coarse-to-fine scale gating ----
+        self.num_scales = num_scales
+        self.c2f_warmup_frac = float(c2f_warmup_frac)
+        # gate per scale (coarse..fine); all-on by default so behaviour is unchanged at frac=0
+        self._scale_gates = [1.0] * num_scales
+        if self.c2f_warmup_frac > 0:
+            self.set_progress(0.0)
+
+    def set_progress(self, frac: float) -> None:
+        """Update coarse-to-fine gates given training progress frac in [0,1].
+        Coarsest scale (index 0) opens first; each successive (finer) scale ramps on
+        later, all fully open by the end of the warmup window. No-op if c2f disabled."""
+        if self.c2f_warmup_frac <= 0:
+            return
+        p = min(1.0, max(0.0, frac / self.c2f_warmup_frac))  # 0..1 over the warmup window
+        S = self.num_scales
+        # gate_i = clamp(p*S - i, 0, 1): scale 0 opens over [0,1/S], scale 1 over [1/S,2/S], ...
+        self._scale_gates = [min(1.0, max(0.0, p * S - i)) for i in range(S)]
 
     # ------------------------------------------------------------------
     # Internal helpers
@@ -478,7 +701,7 @@ def _build_sigma_net(
             nn.init.normal_(out.weight, std=0.01)
             nn.init.zeros_(out.bias)
             layers.append(out)
-            self.sigma_net = nn.Sequential(*layers)
+            self.sigma_net = FusedHiddenMLP(*layers)
         else:
             # Single-linear "explicit" decoder. Small init -> density ~ 0 initially.
             self.sigma_net = nn.Linear(self.feature_dim, 1, bias=True)
@@ -491,14 +714,25 @@ def _build_sigma_net(
 
     def get_densities(self, coords: torch.Tensor) -> torch.Tensor:
         pts = coords.reshape(-1, 3)
-        R = self.so3.as_matrix()  # (T, 3, 3)
+        R = self._rotation_matrices_override
+        if R is None:
+            R = self.so3.as_matrix()  # (T, 3, 3)
+        gates = self._scale_gates if getattr(self, "c2f_warmup_frac", 0.0) > 0 else None
         features = interpolate_ms_features_tilted(
             pts=pts,
             ms_grids=self.grids,
             rotation_matrices=R,
+            scale_gates=gates,
+            include_plane_tv=getattr(self, "_plane_tv_fusion_requested", False),
         )
+        plane_tv = None
+        if isinstance(features, tuple):
+            features, plane_tv = features
         density_before_activation = self.sigma_net(features)
-        return self.density_activation(density_before_activation)
+        density = self.density_activation(density_before_activation)
+        if plane_tv is not None:
+            return density, plane_tv
+        return density
 
     def forward(self, pts: torch.Tensor) -> torch.Tensor:
         return self.get_densities(pts)
@@ -604,33 +838,29 @@ def interpolate_ms_features_cp_tilted(
     # Rotate all points by all rotations: (T, B, 3)
     rotated = torch.einsum("tij,bj->tbi", rotation_matrices, pts)
 
+    # For each transform t, we need three 1D samples: at x_t, y_t, z_t.
+    # Lay them out as (3T, B) coords, matching each line grid's first dim.
+    # Axis order per transform: x, y, z.
+    coords_1d = rotated.reshape(T, B, 3).permute(0, 2, 1).reshape(3 * T, B)
+
+    # grid_sample wants 4D input for 2D sampling: sample the (3T, C, 1, L) lines
+    # with 2D coords whose y is fixed at 0. The grid only depends on the points,
+    # so it is built once here rather than per scale inside the loop.
+    grid = torch.stack(
+        [
+            coords_1d,  # x
+            torch.zeros_like(coords_1d),  # y
+        ],
+        dim=-1,
+    ).unsqueeze(1)  # (3T, 1, B, 2)
+
     per_scale_features = []
     for line_coef in ms_grids:
         # line_coef: (3T, C, L)  — three 1D feature lines per transform (x, y, z)
-        C, _ = line_coef.shape[1], line_coef.shape[2]
-
-        # For each transform t, we need three 1D samples: at x_t, y_t, z_t.
-        # Lay them out as (3T, B) coords, matching line_coef's first dim.
-        # Axis order per transform: x, y, z.
-        coords_1d = rotated.reshape(T, B, 3).permute(0, 2, 1).reshape(3 * T, B)
-        # coords_1d: (3T, B), each row is samples along one axis for one transform
-
-        # grid_sample wants 4D input for 2D sampling, or we can use 1D via a
-        # (3T, C, 1, L) reshape and pass 2D coords with y fixed at 0.
-        # Simpler: use F.grid_sample with a 4D trick, or just do manual linear interp.
-        # Here's the grid_sample way:
-        line_coef_4d = line_coef.unsqueeze(2)  # (3T, C, 1, L)
-        # grid: need (3T, Hout=1, Wout=B, 2), with x = coord, y = 0
-        grid = torch.stack(
-            [
-                coords_1d,  # x
-                torch.zeros_like(coords_1d),  # y
-            ],
-            dim=-1,
-        ).unsqueeze(1)  # (3T, 1, B, 2)
+        C = line_coef.shape[1]
 
         sampled = F.grid_sample(
-            line_coef_4d,
+            line_coef.unsqueeze(2),  # (3T, C, 1, L)
             grid,
             align_corners=True,
             mode="bilinear",
diff --git a/src/quantem/core/ml/models/so3params.py b/src/quantem/core/ml/models/so3params.py
index cd55d0ac6..1efceff30 100644
--- a/src/quantem/core/ml/models/so3params.py
+++ b/src/quantem/core/ml/models/so3params.py
@@ -42,9 +42,15 @@ def quat_to_rotmat(q: torch.Tensor) -> torch.Tensor:
         wx, wy, wz = w * x, w * y, w * z
         R = torch.stack(
             [
-                1 - 2 * (yy + zz), 2 * (xy - wz),     2 * (xz + wy),
-                2 * (xy + wz),     1 - 2 * (xx + zz), 2 * (yz - wx),
-                2 * (xz - wy),     2 * (yz + wx),     1 - 2 * (xx + yy),
+                1 - 2 * (yy + zz),
+                2 * (xy - wz),
+                2 * (xz + wy),
+                2 * (xy + wz),
+                1 - 2 * (xx + zz),
+                2 * (yz - wx),
+                2 * (xz - wy),
+                2 * (yz + wx),
+                1 - 2 * (xx + yy),
             ],
             dim=-1,
         ).reshape(*q.shape[:-1], 3, 3)
@@ -79,10 +85,18 @@ def rotmat_to_quat(R: torch.Tensor) -> torch.Tensor:
         S0, S1, S2, S3 = S.unbind(-1)
 
         # each candidate in [x, y, z, w] order
-        cand_w = torch.stack([(m21 - m12) / S0, (m02 - m20) / S0, (m10 - m01) / S0, 0.25 * S0], dim=-1)
-        cand_x = torch.stack([0.25 * S1, (m01 + m10) / S1, (m02 + m20) / S1, (m21 - m12) / S1], dim=-1)
-        cand_y = torch.stack([(m01 + m10) / S2, 0.25 * S2, (m12 + m21) / S2, (m02 - m20) / S2], dim=-1)
-        cand_z = torch.stack([(m02 + m20) / S3, (m12 + m21) / S3, 0.25 * S3, (m10 - m01) / S3], dim=-1)
+        cand_w = torch.stack(
+            [(m21 - m12) / S0, (m02 - m20) / S0, (m10 - m01) / S0, 0.25 * S0], dim=-1
+        )
+        cand_x = torch.stack(
+            [0.25 * S1, (m01 + m10) / S1, (m02 + m20) / S1, (m21 - m12) / S1], dim=-1
+        )
+        cand_y = torch.stack(
+            [(m01 + m10) / S2, 0.25 * S2, (m12 + m21) / S2, (m02 - m20) / S2], dim=-1
+        )
+        cand_z = torch.stack(
+            [(m02 + m20) / S3, (m12 + m21) / S3, 0.25 * S3, (m10 - m01) / S3], dim=-1
+        )
 
         cands = torch.stack([cand_w, cand_x, cand_y, cand_z], dim=-2)  # (..., 4, 4)
         idx = t.argmax(dim=-1)  # (...,)
@@ -174,11 +188,13 @@ def rotmat_to_r9(R: torch.Tensor) -> torch.Tensor:
     @staticmethod
     def r9_to_rotmat(M: torch.Tensor) -> torch.Tensor:
         """R9 (..., 3, 3) -> nearest SO(3) matrix via SVD+."""
-        U, _, Vh = torch.linalg.svd(M)
-        d = torch.det(U @ Vh)
-        diag = torch.ones(*M.shape[:-2], 3, device=M.device, dtype=M.dtype)
-        diag[..., 2] = d
-        return U @ (diag.unsqueeze(-1) * Vh)
+        # Precision-critical rotations stay fp32 under autocast, preserving fused dispatch.
+        with torch.autocast(device_type=M.device.type, enabled=False):
+            U, _, Vh = torch.linalg.svd(M)
+            d = torch.det(U @ Vh)
+            diag = torch.ones(*M.shape[:-2], 3, device=M.device, dtype=M.dtype)
+            diag[..., 2] = d
+            return U @ (diag.unsqueeze(-1) * Vh)
 
     def as_matrix(self) -> torch.Tensor:
         return self.r9_to_rotmat(self.M)
@@ -189,4 +205,4 @@ def from_matrix(cls, R: torch.Tensor) -> "SO3ParamR9SVD":
         obj = cls(R.shape[0], init="identity")
         with torch.no_grad():
             obj.M.copy_(cls.rotmat_to_r9(R))
-        return obj
\ No newline at end of file
+        return obj
diff --git a/src/quantem/core/ml/optimizer_mixin.py b/src/quantem/core/ml/optimizer_mixin.py
index ece31dff9..a32dc6895 100644
--- a/src/quantem/core/ml/optimizer_mixin.py
+++ b/src/quantem/core/ml/optimizer_mixin.py
@@ -294,14 +294,15 @@ class Plateau:
         _name: str = "plateau"
 
         def params(self, base_LR: float, num_iter: int | None = None) -> dict:
-            if self.min_lr is None:
-                self.min_lr = self.min_lr_factor * base_LR
+            # Derived values stay local: params() must not mutate the dataclass, or a
+            # shared instance bakes in the first optimizer's base LR for every later one.
+            min_lr = self.min_lr if self.min_lr is not None else self.min_lr_factor * base_LR
             return {
                 "mode": self.mode,
                 "factor": self.factor,
                 "patience": self.patience,
                 "threshold": self.threshold,
-                "min_lr": self.min_lr,
+                "min_lr": min_lr,
                 "cooldown": self.cooldown,
             }
 
@@ -332,13 +333,12 @@ def params(self, base_LR: float, num_iter: int | None = None) -> dict:
             if effective_num_iter is None:
                 raise ValueError("num_iter must be set if num_iter is not provided")
 
-            self.num_iter = effective_num_iter
-
+            gamma = self.gamma
             if self.factor is not None:
-                self.gamma = self.factor ** (1.0 / effective_num_iter)
+                gamma = self.factor ** (1.0 / effective_num_iter)
 
             return {
-                "gamma": self.gamma,
+                "gamma": gamma,
             }
 
     @dataclass
@@ -385,13 +385,11 @@ class Cyclic:
         _name: str = "cyclic"
 
         def params(self, base_LR: float, num_iter: int | None = None) -> dict:
-            if self.base_lr is None:
-                self.base_lr = self.base_lr_factor * base_LR
-            if self.max_lr is None:
-                self.max_lr = self.max_lr_factor * base_LR
+            base_lr = self.base_lr if self.base_lr is not None else self.base_lr_factor * base_LR
+            max_lr = self.max_lr if self.max_lr is not None else self.max_lr_factor * base_LR
             return {
-                "base_lr": self.base_lr,
-                "max_lr": self.max_lr,
+                "base_lr": base_lr,
+                "max_lr": max_lr,
                 "step_size_up": self.step_size_up,
                 "step_size_down": self.step_size_down,
                 "mode": self.mode,
@@ -426,12 +424,11 @@ def params(self, base_LR: float, num_iter: int | None = None) -> dict:
                 raise ValueError(
                     "total_iters must be set if num_iter is not provided"
                 )  # Should never be reached
-            if self.total_iters is None:
-                self.total_iters = num_iter
+            total_iters = self.total_iters if self.total_iters is not None else num_iter
             return {
                 "start_factor": self.start_factor,
                 "end_factor": self.end_factor,
-                "total_iters": self.total_iters,
+                "total_iters": total_iters,
             }
 
     @dataclass
@@ -459,10 +456,9 @@ def params(self, base_LR: float, num_iter: int | None = None) -> dict:
                 raise ValueError(
                     "T_max must be set if num_iter is not provided"
                 )  # Should never be reached
-            if self.T_max is None:
-                self.T_max = num_iter
+            T_max = self.T_max if self.T_max is not None else num_iter
             return {
-                "T_max": self.T_max,
+                "T_max": T_max,
                 "eta_min": self.eta_min,
             }
 
@@ -572,7 +568,9 @@ def _normalize_optimizer_params(
         if isinstance(params, OptimizerParamsType):
             return {self.DEFAULT_OPTIMIZER_KEY: params}
         if not isinstance(params, dict):
-            raise TypeError(f"optimizer_params must be OptimizerParamsType or dict, got {type(params)}")
+            raise TypeError(
+                f"optimizer_params must be OptimizerParamsType or dict, got {type(params)}"
+            )
         # Single optimizer as dict shorthand, e.g. {"name": "adam", "lr": 1e-3}
         if self._is_single_optimizer_dict(params):
             return {self.DEFAULT_OPTIMIZER_KEY: OptimizerParams.parse_dict(d=params)}
@@ -597,7 +595,9 @@ def scheduler_params(self, params: SchedulerParamsType | dict):
         if isinstance(params, dict):
             params = SchedulerParams.parse_dict(d=params)
         if not isinstance(params, SchedulerParamsType):
-            raise TypeError(f"scheduler parameters must be a SchedulerParamsType, got {type(params)}")
+            raise TypeError(
+                f"scheduler parameters must be a SchedulerParamsType, got {type(params)}"
+            )
         self._scheduler_params = params
 
     @abstractmethod
@@ -662,7 +662,9 @@ def set_optimizer(self, opt_params: OptimizerParamsType | dict | None = None) ->
         for key, tensors in groups.items():
             for p in tensors:
                 p.requires_grad_(True)
-            param_groups.append({"params": tensors, **specs[key].params()})
+            # "name" records the group key so reconnect_optimizer_to_parameters
+            # can re-join hyperparameters to groups by key, not position.
+            param_groups.append({"params": tensors, **specs[key].params(), "name": key})
         self._optimizer = self._build_optimizer(spec_list[0], param_groups)
 
     def _build_optimizer(self, opt_params, param_groups) -> "torch.optim.Optimizer":
@@ -672,11 +674,32 @@ def _build_optimizer(self, opt_params, param_groups) -> "torch.optim.Optimizer":
         so each group's ``lr`` etc. overrides the optimizer-level default. ``NoneOptimizer`` must
         have been filtered out by the caller.
         """
+        # Fused Adam/AdamW runs the whole step in one CUDA kernel (~2x faster than the
+        # default foreach path on large grids, same update rule); only valid when every
+        # parameter lives on a CUDA device.
+        fused = all(p.is_cuda for group in param_groups for p in group["params"])
+        cuda_graphs = (
+            bool(getattr(self, "_cuda_graphs_optimizer", False))
+            and fused
+            and isinstance(opt_params, OptimizerParams.Adam)
+        )
+        if cuda_graphs:
+            # Captured Adam must read its learning rate from stable device storage so
+            # scheduler updates made between replays remain visible to the graph. This
+            # relies on PyTorch 2.12 schedulers updating tensor LRs in place via ``fill_``;
+            # older PyTorch versions may replace the value and silently freeze the LR
+            # seen by an already-captured graph.
+            for group in param_groups:
+                group["lr"] = torch.tensor(group["lr"], device=group["params"][0].device)
         match opt_params:
             case OptimizerParams.Adam():
-                return torch.optim.Adam(param_groups)
+                return torch.optim.Adam(
+                    param_groups,
+                    fused=fused,
+                    capturable=cuda_graphs,
+                )
             case OptimizerParams.AdamW():
-                return torch.optim.AdamW(param_groups)
+                return torch.optim.AdamW(param_groups, fused=fused)
             case OptimizerParams.SGD():
                 return torch.optim.SGD(param_groups)
             case OptimizerParams.NoneOptimizer():
@@ -688,7 +711,9 @@ def _build_optimizer(self, opt_params, param_groups) -> "torch.optim.Optimizer":
                 raise NotImplementedError(f"Unknown optimizer type: {opt_params}")
 
     def set_scheduler(
-        self, scheduler_params: SchedulerParamsType | dict | None = None, num_iter: int | None = None
+        self,
+        scheduler_params: SchedulerParamsType | dict | None = None,
+        num_iter: int | None = None,
     ) -> None:
         """Set the scheduler for this model."""
         if scheduler_params is not None:
@@ -761,7 +786,7 @@ def has_optimizer(self) -> bool:
     def get_current_lr(self) -> float:
         """Get the current learning rate."""
         if self._optimizer is not None:
-            return self._optimizer.param_groups[0]["lr"]
+            return float(self._optimizer.param_groups[0]["lr"])
         return 0.0
 
     def remove_optimizer(self) -> None:
@@ -802,14 +827,26 @@ def reconnect_optimizer_to_parameters(self) -> None:
         old_hyperparams = [
             {k: v for k, v in pg.items() if k != "params"} for pg in self._optimizer.param_groups
         ]
+        old_by_name = {hp["name"]: hp for hp in old_hyperparams if "name" in hp}
 
         self._optimizer.param_groups.clear()
-        for tensors in new_groups.values():
-            self._optimizer.add_param_group({"params": tensors})
-
-        # Restore per-group hyperparameters by index
-        for new_pg, old_pg in zip(self._optimizer.param_groups, old_hyperparams):
-            new_pg.update(old_pg)
+        for key, tensors in new_groups.items():
+            self._optimizer.add_param_group({"params": tensors, "name": key})
+
+        if old_by_name:
+            # Re-join hyperparameters to groups by key: group order/membership
+            # from get_optimization_parameters() is not contractual, and index
+            # alignment silently attaches the wrong lr when it changes. Groups
+            # with no old counterpart keep the optimizer defaults.
+            for new_pg in self._optimizer.param_groups:
+                hp = old_by_name.get(new_pg["name"])
+                if hp is not None:
+                    new_pg.update(hp)
+        else:
+            # Optimizers restored from checkpoints predating group names:
+            # index alignment is the only association available.
+            for new_pg, old_pg in zip(self._optimizer.param_groups, old_hyperparams):
+                new_pg.update(old_pg)
 
         # Remap state for tensors that survived
         new_state = {}
diff --git a/src/quantem/core/ml/profiling.py b/src/quantem/core/ml/profiling.py
index aec6e5edb..02757b13c 100644
--- a/src/quantem/core/ml/profiling.py
+++ b/src/quantem/core/ml/profiling.py
@@ -1,3 +1,4 @@
+import os
 from contextlib import contextmanager
 
 import torch.cuda.nvtx as nvtx
@@ -12,3 +13,50 @@ def nvtx_range(enabled: bool, name: str):
     finally:
         if enabled:
             nvtx.range_pop()
+
+
+# --- Bounded nsys capture (project-scoped, bench-profiling lineage) -----------------
+#
+# When QUANTEM_NSYS_CAPTURE="<start_step>:<n_steps>" is set (e.g. "10:5"), every rank
+# calls torch.cuda.profiler.start() at grad step <start_step> and .stop() after
+# <n_steps> more steps. Paired with
+#   nsys profile --capture-range=cudaProfilerApi --capture-range-end=stop
+# this bounds the trace to a few steady-state iterations so per-rank reports stay
+# small enough to open. Inert (single early return) when the env var is unset.
+
+_NSYS_SPEC: tuple[int, int] | None = None
+_NSYS_STATE: str = "unparsed"  # unparsed -> idle -> started -> stopped / disabled
+
+
+def nsys_capture_tick(grad_step: int) -> None:
+    global _NSYS_SPEC, _NSYS_STATE
+    if _NSYS_STATE in ("stopped", "disabled"):
+        return
+    if _NSYS_STATE == "unparsed":
+        raw = os.environ.get("QUANTEM_NSYS_CAPTURE", "")
+        if not raw:
+            _NSYS_STATE = "disabled"
+            return
+        try:
+            start_s, n_s = raw.split(":")
+            start, n = int(start_s), int(n_s)
+            if start < 1 or n < 1:
+                raise ValueError
+        except ValueError:
+            print(f"QUANTEM_NSYS_CAPTURE malformed ({raw!r}); expected 'start:n' -> disabled")
+            _NSYS_STATE = "disabled"
+            return
+        _NSYS_SPEC = (start, n)
+        _NSYS_STATE = "idle"
+    assert _NSYS_SPEC is not None
+    start, n = _NSYS_SPEC
+    if _NSYS_STATE == "idle" and grad_step >= start:
+        import torch.cuda.profiler as _prof
+
+        _prof.start()
+        _NSYS_STATE = "started"
+    elif _NSYS_STATE == "started" and grad_step >= start + n:
+        import torch.cuda.profiler as _prof
+
+        _prof.stop()
+        _NSYS_STATE = "stopped"
diff --git a/src/quantem/core/ml/s3im.py b/src/quantem/core/ml/s3im.py
new file mode 100644
index 000000000..858f5ee36
--- /dev/null
+++ b/src/quantem/core/ml/s3im.py
@@ -0,0 +1,214 @@
+from math import exp, isqrt
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def _gaussian(window_size: int, sigma: float) -> torch.Tensor:
+    center = window_size // 2
+    values = [exp(-((x - center) ** 2) / float(2 * sigma**2)) for x in range(window_size)]
+    kernel = torch.tensor(values, dtype=torch.float32)
+    return kernel / kernel.sum()
+
+
+def _create_window(window_size: int, channel: int, *, device, dtype) -> torch.Tensor:
+    window_1d = _gaussian(window_size, 1.5).to(device=device, dtype=dtype).unsqueeze(1)
+    window_2d = window_1d @ window_1d.t()
+    return window_2d.unsqueeze(0).unsqueeze(0).expand(channel, 1, window_size, window_size).contiguous()
+
+
+def _auto_patch_shape(num_pixels: int) -> tuple[int, int]:
+    if num_pixels <= 0:
+        raise ValueError(f"num_pixels must be >= 1, got {num_pixels}")
+
+    root = isqrt(num_pixels)
+    for height in range(root, 0, -1):
+        if num_pixels % height == 0:
+            return height, num_pixels // height
+    return 1, num_pixels
+
+
+class SSIM(nn.Module):
+    def __init__(
+        self,
+        *,
+        window_size: int = 4,
+        stride: int = 4,
+        value_range: float = 1.0,
+        k1: float = 0.01,
+        k2: float = 0.03,
+    ):
+        super().__init__()
+        self.window_size = int(window_size)
+        self.stride = int(stride)
+        self.value_range = float(value_range)
+        self.k1 = float(k1)
+        self.k2 = float(k2)
+        self._cached_channel = 0
+        self._cached_window: torch.Tensor | None = None
+
+        if self.window_size <= 0:
+            raise ValueError(f"window_size must be >= 1, got {self.window_size}")
+        if self.stride <= 0:
+            raise ValueError(f"stride must be >= 1, got {self.stride}")
+        if self.value_range <= 0.0:
+            raise ValueError(f"value_range must be > 0, got {self.value_range}")
+
+    def _window(self, x: torch.Tensor) -> torch.Tensor:
+        channel = int(x.shape[1])
+        if (
+            self._cached_window is None
+            or self._cached_channel != channel
+            or self._cached_window.device != x.device
+            or self._cached_window.dtype != x.dtype
+        ):
+            self._cached_window = _create_window(
+                self.window_size,
+                channel,
+                device=x.device,
+                dtype=x.dtype,
+            )
+            self._cached_channel = channel
+        return self._cached_window
+
+    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+        if x.shape != y.shape:
+            raise ValueError(f"Expected matching shapes, got {tuple(x.shape)} and {tuple(y.shape)}")
+        if x.ndim != 4:
+            raise ValueError(f"Expected [B, C, H, W], got {tuple(x.shape)}")
+
+        window = self._window(x)
+        channel = int(x.shape[1])
+        padding = (self.window_size - 1) // 2
+
+        mu_x = F.conv2d(x, window, padding=padding, groups=channel, stride=self.stride)
+        mu_y = F.conv2d(y, window, padding=padding, groups=channel, stride=self.stride)
+
+        mu_x_sq = mu_x.square()
+        mu_y_sq = mu_y.square()
+        mu_xy = mu_x * mu_y
+
+        sigma_x_sq = (
+            F.conv2d(x * x, window, padding=padding, groups=channel, stride=self.stride) - mu_x_sq
+        )
+        sigma_y_sq = (
+            F.conv2d(y * y, window, padding=padding, groups=channel, stride=self.stride) - mu_y_sq
+        )
+        sigma_xy = (
+            F.conv2d(x * y, window, padding=padding, groups=channel, stride=self.stride) - mu_xy
+        )
+
+        sigma_x_sq = sigma_x_sq.clamp_min(0.0)
+        sigma_y_sq = sigma_y_sq.clamp_min(0.0)
+
+        c1 = (self.k1 * self.value_range) ** 2
+        c2 = (self.k2 * self.value_range) ** 2
+
+        ssim_map = ((2.0 * mu_xy + c1) * (2.0 * sigma_xy + c2)) / (
+            (mu_x_sq + mu_y_sq + c1) * (sigma_x_sq + sigma_y_sq + c2)
+        ).clamp_min(1.0e-12)
+        ssim_map = ssim_map.clamp(min=-1.0, max=1.0)
+        return ssim_map.mean(dim=(1, 2, 3))
+
+
+class S3IMLoss(nn.Module):
+    """Paper-faithful S3IM for scalar or multi-channel tomography ray batches."""
+
+    def __init__(
+        self,
+        *,
+        kernel_size: int = 4,
+        stride: int | None = None,
+        repeat_time: int = 10,
+        patch_height: int | None = None,
+        patch_width: int | None = None,
+        value_range: float = 1.0,
+    ):
+        super().__init__()
+        self.kernel_size = int(kernel_size)
+        self.stride = int(stride) if stride is not None else self.kernel_size
+        self.repeat_time = int(repeat_time)
+        self.patch_height = None if patch_height is None else int(patch_height)
+        self.patch_width = None if patch_width is None else int(patch_width)
+        self.value_range = float(value_range)
+        self.ssim = SSIM(
+            window_size=self.kernel_size,
+            stride=self.stride,
+            value_range=self.value_range,
+        )
+
+        if self.repeat_time <= 0:
+            raise ValueError(f"repeat_time must be >= 1, got {self.repeat_time}")
+        if self.patch_height is not None and self.patch_height <= 0:
+            raise ValueError(f"patch_height must be >= 1, got {self.patch_height}")
+        if self.patch_width is not None and self.patch_width <= 0:
+            raise ValueError(f"patch_width must be >= 1, got {self.patch_width}")
+
+    @staticmethod
+    def _canonicalize(values: torch.Tensor) -> torch.Tensor:
+        if values.ndim == 1:
+            values = values.unsqueeze(-1)
+        if values.ndim != 2:
+            raise ValueError(
+                "Expected batched pixel values shaped [B] or [B, C], "
+                f"got {tuple(values.shape)}"
+            )
+        return values.float()
+
+    def resolve_patch_shape(self, batch_size: int) -> tuple[int, int]:
+        if self.patch_height is None and self.patch_width is None:
+            return _auto_patch_shape(batch_size)
+        if self.patch_height is None or self.patch_width is None:
+            raise ValueError("patch_height and patch_width must both be set or both be omitted.")
+        return self.patch_height, self.patch_width
+
+    def forward(
+        self,
+        predicted: torch.Tensor,
+        target: torch.Tensor,
+        *,
+        return_similarity: bool = False,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        predicted = self._canonicalize(predicted)
+        target = self._canonicalize(target)
+        if predicted.shape != target.shape:
+            raise ValueError(
+                f"predicted and target must have matching shapes, got "
+                f"{tuple(predicted.shape)} and {tuple(target.shape)}"
+            )
+
+        batch_size, channels = predicted.shape
+        patch_height, patch_width = self.resolve_patch_shape(batch_size)
+        expected_size = patch_height * patch_width
+        if expected_size != batch_size:
+            raise ValueError(
+                "S3IM expects patch_height * patch_width == batch size, "
+                f"got {patch_height} * {patch_width} != {batch_size}"
+            )
+
+        index_list = [torch.arange(batch_size, device=predicted.device)]
+        for _ in range(1, self.repeat_time):
+            index_list.append(torch.randperm(batch_size, device=predicted.device))
+        res_index = torch.cat(index_list, dim=0)
+
+        pred_all = predicted[res_index]
+        target_all = target[res_index]
+        pred_patch = pred_all.transpose(0, 1).reshape(
+            1,
+            channels,
+            patch_height,
+            patch_width * self.repeat_time,
+        )
+        target_patch = target_all.transpose(0, 1).reshape(
+            1,
+            channels,
+            patch_height,
+            patch_width * self.repeat_time,
+        )
+
+        similarity = self.ssim(pred_patch, target_patch).mean().clamp(min=-1.0, max=1.0)
+        loss = (1.0 - similarity).clamp(min=0.0, max=2.0)
+        if return_similarity:
+            return loss, similarity
+        return loss
diff --git a/src/quantem/core/quantem.yaml b/src/quantem/core/quantem.yaml
index 80f34587a..0e4913a4f 100644
--- a/src/quantem/core/quantem.yaml
+++ b/src/quantem/core/quantem.yaml
@@ -14,6 +14,10 @@ dtype_complex: complex64
 # global verbosity, not currently used
 
 verbose: 1
+
+# Dispatch to the fused quantem-cuda kernels when installed (pip install quantem[cuda])
+# and the tensors are on a CUDA device. Set false to force the pure-torch paths.
+use_cuda_kernels: true
 cupy:
   # The size of the fft cache in MB used by cupy
   # https://docs.cupy.dev/en/stable/user_guide/fft.html#fft-plan-cache
diff --git a/src/quantem/diffractive_imaging/object_models.py b/src/quantem/diffractive_imaging/object_models.py
index d311d8ddb..17c4e72c1 100644
--- a/src/quantem/diffractive_imaging/object_models.py
+++ b/src/quantem/diffractive_imaging/object_models.py
@@ -396,6 +396,30 @@ def get_tv_loss(
         return loss
 
     def _calc_tv_loss(self, array: torch.Tensor, weight: tuple[float, float]) -> torch.Tensor:
+        # Identical math via the fused quantem-cuda L1 kernel when available.
+        # Weighted size-1 axes fall through to torch so degenerate inputs
+        # behave exactly as before.
+        if (
+            array.ndim == 3
+            and array.is_cuda
+            and array.dtype == torch.float32
+            and config.get("has_quantem_cuda")
+            and config.get("use_cuda_kernels", default=True)
+            and not (weight[0] > 0 and array.shape[0] == 1)
+            and not (weight[1] > 0 and (array.shape[1] == 1 or array.shape[2] == 1))
+        ):
+            from quantem.cuda.core import tv_loss_l1_3d
+
+            s, h, w_ = array.shape
+            wts = array.new_tensor([weight[0], weight[1], weight[1]])
+            active = int((wts > 0).sum())
+            if active == 0:
+                return self._get_zero_loss_tensor()
+            counts = array.new_tensor(
+                [(s - 1) * h * w_, s * (h - 1) * w_, s * h * (w_ - 1)]
+            ).clamp(min=1)
+            return (wts * tv_loss_l1_3d(array) / counts).sum() / active
+
         loss = self._get_zero_loss_tensor()
         calc_dim = 0
         for dim in range(array.ndim):
diff --git a/src/quantem/tomography/__init__.py b/src/quantem/tomography/__init__.py
index e69de29bb..2ec8d60c0 100644
--- a/src/quantem/tomography/__init__.py
+++ b/src/quantem/tomography/__init__.py
@@ -0,0 +1 @@
+from quantem.tomography.atom_trace import Atoms as Atoms
diff --git a/src/quantem/tomography/atom_trace.py b/src/quantem/tomography/atom_trace.py
new file mode 100644
index 000000000..4c4be7f07
--- /dev/null
+++ b/src/quantem/tomography/atom_trace.py
@@ -0,0 +1,1173 @@
+"""Atom tracing from 3D tomographic volumes via differentiable Gaussian splatting.
+
+This module decomposes a reconstructed 3D volume (``Dataset3d``) into a set of
+atomic sites, each modeled as a 3D Gaussian.  Rather than fitting one atom at a
+time (the classic Levenberg-Marquardt approach), the whole volume is treated as a
+single differentiable model and *all* parameters are optimized jointly with Adam.
+Overlap between neighboring atoms is handled automatically by the joint fit.
+
+Forward model
+-------------
+The volume is decomposed into a sharp atomic part and a smooth background, both
+non-negative::
+
+    volume = volume_atoms + volume_background,   volume_atoms >= 0,  volume_background >= 0
+
+``volume_atoms``
+    A sum of 3D Gaussians.  Intensities are kept >= 0, so the atomic part is >= 0.
+``volume_background``
+    A low-degree tensor-product Bezier (Bernstein) field over a small control
+    lattice.  With non-negative control points it is non-negative *everywhere*
+    (Bernstein partition-of-unity) and slowly varying by construction, so it
+    cannot absorb sharp atomic features -- the model class itself separates the
+    two.  This is the hook for future background regularization.
+
+Atom models
+-----------
+isotropic
+    5 degrees of freedom per atom: ``x, y, z`` position, intensity ``I``, and a
+    single width ``sigma``.
+anisotropic (planned)
+    10 degrees of freedom: ``x, y, z``, ``I`` and a Cholesky-parameterized
+    precision matrix ``Lambda = L @ L.T`` (positive-definite by construction, so
+    the Gaussian can never diverge).
+
+Efficiency
+----------
+The renderer never evaluates every Gaussian over every voxel.  Each Gaussian only
+contributes to a small ``(2*window_radius + 1)**3`` window around its rounded
+center, accumulated with ``scatter_add``.  Cost is ``O(n_atoms * window_volume)``
+rather than ``O(n_atoms * n_voxels)``.  A window of 3-4 sigma is the right
+accuracy/speed trade-off for tracing (relative truncation error ~1e-3 at 3 sigma,
+~6e-6 at 4 sigma); larger windows are only needed to *measure* truncation.
+
+Conventions
+-----------
+All internal site coordinates are in **voxel / array-index units** matching the
+volume's axes (axis 0, 1, 2).  Conversion to physical units happens only at the
+public ``sites`` (Vector) boundary, using the volume's ``sampling``/``origin``.
+
+.. note::
+    Under construction. The differentiable forward model below (atom renderer +
+    Bezier background) is complete and verified; the ``Atoms`` class, seeding,
+    schedule, and regularizers follow.
+"""
+
+from __future__ import annotations
+
+import math
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+from tqdm.auto import tqdm
+
+from quantem.core import config
+from quantem.core.datastructures import Dataset3d, Vector
+from quantem.core.io.serialize import AutoSerialize
+
+__all__ = [
+    "Atoms",
+    "render_isotropic",
+    "bernstein_basis",
+    "render_background",
+    "gaussian_blur3d",
+    "seed_peaks",
+    "estimate_nn_spacing",
+]
+
+
+def render_isotropic(
+    positions: Tensor,
+    intensities: Tensor,
+    sigmas: Tensor | float,
+    volume_shape: tuple[int, int, int],
+    window_radius: int,
+) -> Tensor:
+    """Render a sum of isotropic 3D Gaussians into a dense volume by local splatting.
+
+    Each Gaussian contributes only to a ``(2*window_radius + 1)**3`` window around
+    its rounded center.  The window *indices* come from ``round(positions)`` (held
+    constant w.r.t. gradients), while the Gaussian is evaluated at the continuous
+    ``positions``, so gradients flow to ``positions``, ``intensities`` and
+    ``sigmas``.  Choose ``window_radius >~ 3-4 * sigma_max`` so that truncation at
+    the window edge is negligible for tracing.
+
+    Parameters
+    ----------
+    positions : Tensor
+        ``(N, 3)`` float tensor of site centers in voxel/array-index coordinates.
+    intensities : Tensor
+        ``(N,)`` float tensor of Gaussian amplitudes (kept >= 0 by the caller).
+    sigmas : Tensor or float
+        Gaussian width(s) in voxels.  Scalar / shape ``(1,)`` broadcasts to all
+        atoms; otherwise shape ``(N,)``.
+    volume_shape : tuple[int, int, int]
+        Output volume shape ``(D0, D1, D2)``.
+    window_radius : int
+        Half-width (in voxels) of the cubic splat window per atom.
+
+    Returns
+    -------
+    Tensor
+        ``(D0, D1, D2)`` rendered atomic volume, differentiable w.r.t.
+        ``positions``, ``intensities`` and ``sigmas``.
+    """
+    device = positions.device
+    dtype = positions.dtype
+    n_atoms = positions.shape[0]
+    d0, d1, d2 = volume_shape
+
+    if n_atoms == 0:
+        return torch.zeros(volume_shape, device=device, dtype=dtype)
+
+    sigmas = torch.as_tensor(sigmas, device=device, dtype=dtype)
+    if sigmas.ndim == 0:
+        sigmas = sigmas.reshape(1)
+    if sigmas.shape[0] == 1:
+        sigmas = sigmas.expand(n_atoms)
+
+    # Integer window centers -- detached so they carry no gradient.
+    centers = torch.round(positions.detach()).long()  # (N, 3)
+
+    # Cubic window offsets, (W**3, 3).
+    rng = torch.arange(-window_radius, window_radius + 1, device=device)
+    o0, o1, o2 = torch.meshgrid(rng, rng, rng, indexing="ij")
+    offsets = torch.stack((o0.reshape(-1), o1.reshape(-1), o2.reshape(-1)), dim=1)
+
+    # Integer voxel coordinates for every (atom, window-voxel): (N, W**3, 3).
+    vox = centers[:, None, :] + offsets[None, :, :]
+
+    # Continuous displacement from the (sub-voxel) atom center to each voxel.
+    diff = vox.to(dtype) - positions[:, None, :]  # (N, W**3, 3)
+    r2 = (diff * diff).sum(dim=-1)  # (N, W**3)
+    val = intensities[:, None] * torch.exp(-0.5 * r2 / (sigmas[:, None] ** 2))
+
+    # Zero out-of-bounds contributions; clamp their indices to a valid slot so the
+    # scatter is safe (those entries add 0.0).
+    shape_t = torch.tensor(volume_shape, device=device)
+    valid = ((vox >= 0) & (vox < shape_t)).all(dim=-1)  # (N, W**3)
+    val = val * valid.to(dtype)
+    vox = torch.minimum(torch.maximum(vox, torch.zeros_like(shape_t)), shape_t - 1)
+
+    flat_idx = (vox[..., 0] * (d1 * d2) + vox[..., 1] * d2 + vox[..., 2]).reshape(-1)
+    out = torch.zeros(d0 * d1 * d2, device=device, dtype=dtype)
+    out = out.scatter_add(0, flat_idx, val.reshape(-1))
+    return out.reshape(volume_shape)
+
+
+def bernstein_basis(
+    num_samples: int,
+    degree: int,
+    device: torch.device | str | None = None,
+    dtype: torch.dtype = torch.float32,
+) -> Tensor:
+    """Bernstein basis matrix sampled on a uniform grid over ``[0, 1]``.
+
+    ``B[s, i] = C(degree, i) * t**i * (1 - t)**(degree - i)`` with
+    ``t = s / (num_samples - 1)``.  Each row sums to 1 (partition of unity), which
+    makes a Bezier field a convex combination of its control points.
+
+    Parameters
+    ----------
+    num_samples : int
+        Number of evenly spaced sample points (the axis length of the volume).
+    degree : int
+        Polynomial degree; the control lattice has ``degree + 1`` points on this
+        axis.  Low degree -> smoother, more slowly varying field.
+    device, dtype
+        Torch device/dtype for the returned matrix.
+
+    Returns
+    -------
+    Tensor
+        ``(num_samples, degree + 1)`` basis matrix.
+    """
+    t = torch.linspace(0.0, 1.0, num_samples, device=device, dtype=dtype)[:, None]
+    i = torch.arange(degree + 1, device=device, dtype=dtype)[None, :]
+    log_binom = (
+        torch.lgamma(torch.tensor(degree + 1.0, device=device, dtype=dtype))
+        - torch.lgamma(i + 1.0)
+        - torch.lgamma(degree - i + 1.0)
+    )
+    # torch.pow gives 0**0 == 1, so the endpoints interpolate the corner controls.
+    return torch.exp(log_binom) * t.pow(i) * (1.0 - t).pow(degree - i)
+
+
+def render_background(
+    control_points: Tensor,
+    bases: tuple[Tensor, Tensor, Tensor],
+) -> Tensor:
+    """Evaluate a tensor-product Bezier (Bernstein) field over a dense volume.
+
+    With non-negative ``control_points`` the field is non-negative everywhere and
+    bounded by ``[control_points.min(), control_points.max()]`` (partition of
+    unity), and is smooth/slowly varying for low control-lattice degree.
+
+    Parameters
+    ----------
+    control_points : Tensor
+        ``(n0 + 1, n1 + 1, n2 + 1)`` control lattice (kept >= 0 by the caller).
+    bases : tuple of Tensor
+        Per-axis Bernstein bases ``(B0, B1, B2)`` from :func:`bernstein_basis`,
+        with shapes ``(D0, n0+1)``, ``(D1, n1+1)``, ``(D2, n2+1)``.
+
+    Returns
+    -------
+    Tensor
+        ``(D0, D1, D2)`` background field, differentiable w.r.t. ``control_points``.
+    """
+    b0, b1, b2 = bases
+    f = torch.einsum("ijk,ai->ajk", control_points, b0)
+    f = torch.einsum("ajk,bj->abk", f, b1)
+    f = torch.einsum("abk,ck->abc", f, b2)
+    return f
+
+
+def gaussian_blur3d(volume: Tensor, sigma: float) -> Tensor:
+    """Separable 3D Gaussian blur with replicate padding.
+
+    Parameters
+    ----------
+    volume : Tensor
+        ``(D0, D1, D2)`` volume.
+    sigma : float
+        Standard deviation in voxels.  ``sigma <= 0`` returns the input unchanged.
+
+    Returns
+    -------
+    Tensor
+        Blurred ``(D0, D1, D2)`` volume.
+    """
+    if sigma <= 0:
+        return volume
+    radius = max(1, int(math.ceil(3.0 * sigma)))
+    x = torch.arange(-radius, radius + 1, device=volume.device, dtype=volume.dtype)
+    kernel = torch.exp(-0.5 * (x / sigma) ** 2)
+    kernel = kernel / kernel.sum()
+    v = volume[None, None]  # (1, 1, D0, D1, D2)
+    for axis in range(3):
+        shape = [1, 1, 1, 1, 1]
+        shape[2 + axis] = kernel.numel()
+        pad = [0, 0, 0, 0, 0, 0]  # F.pad order is (W_lo, W_hi, H_lo, H_hi, D_lo, D_hi)
+        pad[(2 - axis) * 2] = radius
+        pad[(2 - axis) * 2 + 1] = radius
+        v = F.pad(v, pad, mode="replicate")
+        v = F.conv3d(v, kernel.reshape(shape))
+    return v[0, 0]
+
+
+def seed_peaks(
+    volume: Tensor,
+    blur_sigmas: tuple[float, float] | float = (1.0, 2.0),
+    threshold: float | None = None,
+    threshold_fraction: float = 0.99,
+    min_distance: float = 0.0,
+    max_peaks: int | None = None,
+    progress: bool = False,
+) -> tuple[Tensor, Tensor]:
+    """Seed candidate atom sites by difference-of-Gaussians peak detection.
+
+    Bandpass-filters the volume (single blur, or a difference of two blurs to also
+    suppress smooth background), finds 3x3x3 local maxima above a threshold,
+    refines each to sub-voxel accuracy with a per-axis parabolic fit, and
+    optionally enforces a minimum spacing (greedy, brightest-first).
+
+    Parameters
+    ----------
+    volume : Tensor
+        ``(D0, D1, D2)`` volume.
+    blur_sigmas : tuple[float, float] or float
+        Two sigmas -> difference-of-Gaussians (bandpass).  One sigma -> single
+        matched-filter blur.  In voxels.
+    threshold : float or None
+        Absolute cutoff on the filtered (difference-of-Gaussians) response: a voxel
+        is a candidate peak only if its response exceeds this.  Overrides
+        ``threshold_fraction``.
+    threshold_fraction : float
+        Detection threshold as a quantile (0-1) of the filtered response, used when
+        ``threshold`` is None.  Only voxels above this quantile are kept, so
+        ``0.99`` keeps the brightest ~1% of the response (more, weaker sites) and
+        ``0.999`` the brightest ~0.1% (fewer, stronger sites).  Robust to outliers
+        and to the absolute data scale.
+    min_distance : float
+        Minimum spacing between seeds in voxels; closer (dimmer) peaks are
+        dropped.  0 disables.
+    max_peaks : int or None
+        Keep at most this many seeds (brightest first).
+    progress : bool
+        Show a tqdm progress bar over the detection stages.
+
+    Returns
+    -------
+    positions : Tensor
+        ``(M, 3)`` sub-voxel seed coordinates in voxel/array-index units.
+    intensities : Tensor
+        ``(M,)`` volume value sampled at each seed (integer voxel).
+    """
+    # Peak detection runs on CPU for cross-device determinism: the local-maximum
+    # test relies on exact float equality (resp == max_pool(resp)), which is not
+    # reliable on MPS. Seeding is fast and one-time, so this is cheap.
+    bar = tqdm(total=4, desc="find_initial", disable=not progress)
+    volume = volume.detach().to("cpu")
+    dtype = volume.dtype
+    if isinstance(blur_sigmas, (tuple, list)):
+        s_lo, s_hi = blur_sigmas
+        resp = gaussian_blur3d(volume, float(s_lo)) - gaussian_blur3d(volume, float(s_hi))
+    else:
+        resp = gaussian_blur3d(volume, float(blur_sigmas))
+    if threshold is None:
+        threshold = float(np.quantile(resp.numpy(), threshold_fraction))
+    bar.set_postfix_str("DoG bandpass")
+    bar.update(1)
+
+    # 3x3x3 local maxima above threshold (max_pool uses -inf padding, so the
+    # equality test is exact for the window argmax).
+    pooled = F.max_pool3d(resp[None, None], kernel_size=3, stride=1, padding=1)[0, 0]
+    is_peak = (resp == pooled) & (resp > threshold)
+    # Drop the outer shell so the parabolic fit always has neighbors.
+    is_peak[[0, -1], :, :] = False
+    is_peak[:, [0, -1], :] = False
+    is_peak[:, :, [0, -1]] = False
+    idx = is_peak.nonzero(as_tuple=False)  # (M, 3) long
+    bar.set_postfix_str(f"{idx.shape[0]} maxima")
+    bar.update(1)
+    if idx.shape[0] == 0:
+        bar.close()
+        empty = torch.zeros((0, 3), dtype=dtype)
+        return empty, empty[:, 0]
+
+    i, j, k = idx[:, 0], idx[:, 1], idx[:, 2]
+    f0 = resp[i, j, k]
+    offsets = torch.zeros_like(idx, dtype=dtype)
+    eps = torch.finfo(dtype).eps
+    neighbor_pairs = (
+        (resp[i - 1, j, k], resp[i + 1, j, k]),
+        (resp[i, j - 1, k], resp[i, j + 1, k]),
+        (resp[i, j, k - 1], resp[i, j, k + 1]),
+    )
+    for axis, (f_lo, f_hi) in enumerate(neighbor_pairs):
+        denom = f_lo - 2.0 * f0 + f_hi
+        shift = torch.where(
+            denom.abs() > eps, 0.5 * (f_lo - f_hi) / denom, torch.zeros_like(denom)
+        )
+        offsets[:, axis] = shift.clamp(-0.5, 0.5)
+
+    positions = idx.to(dtype) + offsets
+    intensities = volume[i, j, k]
+    # Sort brightest-first.
+    order = torch.argsort(intensities, descending=True)
+    positions, intensities = positions[order], intensities[order]
+    bar.set_postfix_str("subpixel refine")
+    bar.update(1)
+
+    # Greedy minimum-distance suppression (brightest wins).
+    if min_distance > 0 and positions.shape[0] > 1:
+        from scipy.spatial import cKDTree
+
+        pts = positions.numpy()
+        neighbors = cKDTree(pts).query_ball_point(pts, r=float(min_distance))
+        suppressed = np.zeros(pts.shape[0], dtype=bool)
+        keep: list[int] = []
+        for a in tqdm(range(pts.shape[0]), desc="dedup", disable=not progress, leave=False):
+            if suppressed[a]:
+                continue
+            keep.append(a)
+            for b in neighbors[a]:
+                if b > a:
+                    suppressed[b] = True
+        keep_t = torch.as_tensor(keep)
+        positions, intensities = positions[keep_t], intensities[keep_t]
+    bar.set_postfix_str(f"{positions.shape[0]} sites")
+    bar.update(1)
+    bar.close()
+
+    if max_peaks is not None and positions.shape[0] > max_peaks:
+        positions, intensities = positions[:max_peaks], intensities[:max_peaks]
+
+    return positions, intensities
+
+
+def estimate_nn_spacing(
+    volume: Tensor,
+    r_min: float = 2.0,
+    r_max: float | None = None,
+    return_profile: bool = False,
+) -> float | tuple[float, np.ndarray, np.ndarray]:
+    """Estimate the nearest-neighbor spacing (in voxels) from the volume autocorrelation.
+
+    The autocorrelation ``IFFT(|FFT(v)|^2)`` is radially averaged; the radius of
+    its first peak beyond the central self-peak is the first coordination shell,
+    i.e. the typical nearest-neighbor spacing.  Using ``~0.75 *`` this value as the
+    seeding ``min_distance`` strongly suppresses duplicate/false-positive sites.
+
+    Parameters
+    ----------
+    volume : Tensor or ndarray
+        3D volume.
+    r_min : float
+        Ignore peaks closer than this radius (excludes the central self-peak).
+    r_max : float or None
+        Largest radius to consider.  Default ``min(shape) // 2``.
+    return_profile : bool
+        If True, also return the radial-autocorrelation profile (normalized so the
+        zero-lag value is 1) for plotting/inspection.
+
+    Returns
+    -------
+    float or tuple
+        The nearest-neighbor spacing in voxels; or, if ``return_profile``,
+        ``(spacing, radii, radial)`` with ``radii``/``radial`` as 1D arrays.
+
+    Raises
+    ------
+    RuntimeError
+        If no autocorrelation shell peak is found.
+    """
+    from scipy.signal import find_peaks
+
+    vt = torch.as_tensor(volume).detach().to(device="cpu", dtype=torch.float32)
+    vt = vt - vt.mean()
+    power = torch.fft.fftn(vt).abs() ** 2
+    ac = torch.fft.fftshift(torch.fft.ifftn(power).real).numpy()
+    shape = ac.shape
+    center = [s // 2 for s in shape]
+    sq = [((np.arange(s) - c) ** 2).astype(np.float32) for s, c in zip(shape, center)]
+    radius = np.sqrt(sq[0][:, None, None] + sq[1][None, :, None] + sq[2][None, None, :])
+    r_int = np.rint(radius).astype(np.int64).ravel()
+    counts = np.bincount(r_int)
+    radial = np.bincount(r_int, weights=ac.ravel().astype(np.float64)) / np.maximum(counts, 1)
+    cutoff = int(r_max) if r_max is not None else min(shape) // 2
+    radial = radial[:cutoff]
+    if radial[0] > 0:
+        radial = radial / radial[0]  # normalize so the zero-lag value is 1
+    peaks, _ = find_peaks(radial)
+    shell = [int(p) for p in peaks if p >= max(1, int(round(r_min)))]
+    if not shell:
+        raise RuntimeError(
+            "Could not estimate NN spacing from the autocorrelation; "
+            "pass min_distance to find_initial() explicitly."
+        )
+    spacing = float(shell[0])
+    if return_profile:
+        return spacing, np.arange(len(radial), dtype=float), radial
+    return spacing
+
+
+def _inverse_softplus(y: Tensor) -> Tensor:
+    """Numerically stable inverse of ``softplus`` (no overflow for large y).
+
+    ``log(exp(y) - 1) = y + log(-expm1(-y))``; for large y the second term -> 0,
+    so the result stays finite (the naive ``log(expm1(y))`` overflows in float32
+    around y ~ 88).
+    """
+    y = y.clamp(min=1e-6)
+    return y + torch.log(-torch.expm1(-y))
+
+
+class Atoms(AutoSerialize):
+    """Atomic sites traced from a 3D volume by differentiable Gaussian splatting.
+
+    The volume is modeled as ``volume_atoms + volume_background`` (both >= 0) and
+    all parameters are optimized jointly.  Atom parameters are stored internally in
+    **voxel/array-index** coordinates; the public ``sites`` Vector reports them in
+    the source volume's physical units.
+
+    Construct with :meth:`from_dataset`, then call :meth:`find_initial` (and,
+    soon, ``trace``/``optimize``) to populate and refine the sites.
+
+    Parameters
+    ----------
+    volume : Dataset3d
+        Input 3D volume (e.g. a tomographic reconstruction).
+    model : {"isotropic"}
+        Atom model.  Only isotropic (x, y, z, intensity, sigma) is implemented;
+        anisotropic is planned.
+    sigma_init : float
+        Initial Gaussian width in voxels.
+    sigma_cutoff : float
+        Splat window half-width in units of sigma (3-4 is the speed/accuracy
+        sweet spot).
+    background : bool
+        If True, model a smooth non-negative Bezier background.
+    background_degree : int or tuple[int, int, int]
+        Per-axis degree of the background control lattice (lower = smoother).
+    device : str, int, or None
+        Compute device; None selects the quantem default (cuda > mps > cpu).
+    """
+
+    _token = object()
+
+    def __init__(
+        self,
+        volume: Dataset3d,
+        *,
+        model: str = "isotropic",
+        sigma_init: float = 2.0,
+        sigma_cutoff: float = 4.0,
+        background: bool = True,
+        background_degree: int | tuple[int, int, int] = 4,
+        device: str | int | None = None,
+        _token: object | None = None,
+    ):
+        if _token is not self._token:
+            raise RuntimeError("Use Atoms.from_dataset() to instantiate this class.")
+        super().__init__()
+        if model != "isotropic":
+            raise NotImplementedError(
+                "Only model='isotropic' is implemented; anisotropic is planned."
+            )
+        dev, _ = config.validate_device(device)
+        self.device = dev
+        self._model = model
+        self._sigma_init = float(sigma_init)
+        if self._sigma_init <= 0:
+            raise ValueError(f"sigma_init must be > 0, got {self._sigma_init}.")
+        self._sigma_cutoff = float(sigma_cutoff)
+        self._dtype = torch.float32
+
+        # Calibration from the source volume (for voxel <-> physical conversion).
+        self._source = volume
+        self._sampling = np.asarray(volume.sampling, dtype=float)
+        self._origin = np.asarray(volume.origin, dtype=float)
+        self._units = list(volume.units)
+        self._signal_units = volume.signal_units
+
+        # Raw volume as a torch tensor on the chosen device.
+        arr = volume.array
+        vol = torch.as_tensor(np.asarray(arr)) if arr is not None else volume.tensor
+        self._volume = vol.to(device=self.device, dtype=self._dtype)
+        self._shape = tuple(int(s) for s in self._volume.shape)
+
+        # Atom parameters (empty until seeded). Stored raw/unconstrained.
+        self._positions = torch.zeros((0, 3), dtype=self._dtype, device=self.device)
+        self._raw_intensity = torch.zeros((0,), dtype=self._dtype, device=self.device)
+        self._raw_sigma = torch.zeros((0,), dtype=self._dtype, device=self.device)
+
+        # Bezier background control lattice + per-axis Bernstein bases.
+        self._has_background = bool(background)
+        if self._has_background:
+            degs = (
+                (background_degree,) * 3
+                if isinstance(background_degree, int)
+                else tuple(int(d) for d in background_degree)
+            )
+            self._bg_degrees = degs
+            self._bases = tuple(
+                bernstein_basis(self._shape[a], degs[a], device=self.device, dtype=self._dtype)
+                for a in range(3)
+            )
+            init_bg = max(float(self._volume.min()), 1e-4)
+            self._raw_background = _inverse_softplus(
+                torch.full(
+                    tuple(d + 1 for d in degs), init_bg, dtype=self._dtype, device=self.device
+                )
+            )
+        else:
+            self._bg_degrees = None
+            self._bases = None
+            self._raw_background = None
+
+    # ------------------------------------------------------------------ #
+    # Construction
+    # ------------------------------------------------------------------ #
+    @classmethod
+    def from_dataset(
+        cls,
+        volume: Dataset3d,
+        *,
+        model: str = "isotropic",
+        sigma_init: float = 2.0,
+        sigma_cutoff: float = 4.0,
+        background: bool = True,
+        background_degree: int | tuple[int, int, int] = 4,
+        device: str | int | None = None,
+    ) -> "Atoms":
+        """Create an :class:`Atoms` tracer from a 3D :class:`Dataset3d` volume."""
+        if not isinstance(volume, Dataset3d):
+            raise TypeError(f"volume must be a Dataset3d, got {type(volume).__name__}.")
+        return cls(
+            volume,
+            model=model,
+            sigma_init=sigma_init,
+            sigma_cutoff=sigma_cutoff,
+            background=background,
+            background_degree=background_degree,
+            device=device,
+            _token=cls._token,
+        )
+
+    # ------------------------------------------------------------------ #
+    # Constrained parameter views (raw -> physical-meaning)
+    # ------------------------------------------------------------------ #
+    @property
+    def _intensity(self) -> Tensor:
+        return F.softplus(self._raw_intensity)
+
+    @property
+    def _sigma(self) -> Tensor:
+        return F.softplus(self._raw_sigma)
+
+    @property
+    def _background_control(self) -> Tensor | None:
+        return None if not self._has_background else F.softplus(self._raw_background)
+
+    def _window_radius(self) -> int:
+        smax = self._sigma_init if self.num_sites == 0 else float(self._sigma.detach().max())
+        return max(1, int(math.ceil(self._sigma_cutoff * smax)))
+
+    # ------------------------------------------------------------------ #
+    # State
+    # ------------------------------------------------------------------ #
+    @property
+    def num_sites(self) -> int:
+        """Number of atomic sites."""
+        return int(self._positions.shape[0])
+
+    @property
+    def model(self) -> str:
+        """Atom model name."""
+        return self._model
+
+    # ------------------------------------------------------------------ #
+    # Seeding
+    # ------------------------------------------------------------------ #
+    def estimate_spacing(
+        self,
+        r_min: float = 2.0,
+        recompute: bool = False,
+        plot: bool = False,
+        returnfig: bool = False,
+    ) -> float | tuple:
+        """Estimate (and cache) the nearest-neighbor atomic spacing, in voxels.
+
+        Uses the volume autocorrelation (:func:`estimate_nn_spacing`).  The result
+        seeds the default ``min_distance`` for :meth:`find_initial`, which strongly
+        suppresses duplicate / false-positive sites (especially around weak ones).
+
+        Parameters
+        ----------
+        r_min : float, default 2.0
+            Ignore autocorrelation peaks closer than this radius (excludes the
+            central self-peak).
+        recompute : bool, default False
+            Recompute even if a cached value exists.
+        plot : bool, default False
+            Plot the radial autocorrelation profile with the detected first-shell
+            peak (= the spacing) and the resulting ``min_distance`` marked, to help
+            you judge whether the estimate is sensible.
+        returnfig : bool, default False
+            If True, return ``(spacing, fig, ax)`` instead of just the spacing.
+
+        Returns
+        -------
+        float or tuple
+            The spacing in voxels, or ``(spacing, fig, ax)`` if ``returnfig``.
+        """
+        if recompute or getattr(self, "_nn_spacing", None) is None:
+            self._nn_spacing, radii, radial = estimate_nn_spacing(
+                self._volume, r_min=r_min, return_profile=True
+            )
+            self._nn_profile = (radii, radial)
+
+        if not plot:
+            return self._nn_spacing
+
+        import matplotlib.pyplot as plt
+
+        radii, radial = self._nn_profile
+        spacing = self._nn_spacing
+        # Show the shell structure beyond the central self-peak.
+        r_hi = min(len(radial) - 1, max(int(round(4 * spacing)), 12))
+        fig, ax = plt.subplots(figsize=(6, 3.2))
+        ax.plot(radii[1 : r_hi + 1], radial[1 : r_hi + 1], color="0.2", lw=1.2)
+        ax.axvline(spacing, color="tab:red", ls="--", label=f"NN spacing = {spacing:.1f} voxels")
+        ax.axvline(
+            0.75 * spacing, color="tab:blue", ls=":",
+            label=f"min_distance = {0.75 * spacing:.2f} voxels",
+        )
+        ax.set_xlabel("radius (voxels)")
+        ax.set_ylabel("radial autocorrelation")
+        ax.set_title("volume autocorrelation")
+        ax.legend(fontsize=9)
+        fig.tight_layout()
+        return (self._nn_spacing, fig, ax) if returnfig else self._nn_spacing
+
+    def find_initial(
+        self,
+        blur_sigmas: tuple[float, float] | float | None = None,
+        threshold: float | None = None,
+        threshold_fraction: float = 0.99,
+        min_distance: float | None = None,
+        spacing_factor: float = 0.75,
+        max_peaks: int | None = None,
+        progress: bool = True,
+    ) -> "Atoms":
+        """Find an initial set of atomic sites by difference-of-Gaussians detection.
+
+        Band-pass filters the volume to enhance atom-sized blobs, keeps the local
+        maxima above a threshold as candidate sites, refines each to sub-voxel
+        accuracy, and drops duplicates closer than ``min_distance``.  This is the
+        first step of tracing and seeds :meth:`refine`.
+
+        Parameters
+        ----------
+        blur_sigmas : tuple[float, float] or float or None
+            Difference-of-Gaussians widths in voxels, ``(small, large)``; the
+            band-pass highlights features at the atom scale.  Default brackets
+            ``sigma_init`` as ``(0.75, 1.5) * sigma_init``.
+        threshold : float or None
+            Absolute cutoff on the filtered response (a voxel is a candidate only
+            if its response exceeds this).  Overrides ``threshold_fraction``.
+        threshold_fraction : float, default 0.99
+            Detection threshold as a quantile (0-1) of the filtered response: only
+            voxels whose response is above this quantile are kept.  ``0.99`` keeps
+            the brightest ~1% of the response (more sites, including weaker ones);
+            ``0.999`` keeps the brightest ~0.1% (fewer, stronger sites).  Being a
+            quantile, it is robust to outliers and to the absolute data scale.
+        min_distance : float or None
+            Minimum spacing between sites in voxels; of two sites closer than this,
+            the dimmer is removed.  Defaults to ``spacing_factor`` times the
+            nearest-neighbor spacing from :meth:`estimate_spacing` -- the main lever
+            against duplicate / false-positive sites.
+        spacing_factor : float, default 0.75
+            Fraction of the estimated nearest-neighbor spacing used for the default
+            ``min_distance`` (ignored when ``min_distance`` is given).
+        max_peaks : int or None
+            If set, keep only the brightest ``max_peaks`` sites.
+        progress : bool, default True
+            Show a tqdm progress bar.
+
+        Returns
+        -------
+        Atoms
+            ``self``, with ``sites`` populated.
+        """
+        if blur_sigmas is None:
+            blur_sigmas = (0.75 * self._sigma_init, 1.5 * self._sigma_init)
+        if min_distance is None:
+            try:
+                min_distance = spacing_factor * self.estimate_spacing()
+            except RuntimeError:
+                min_distance = 2.0 * self._sigma_init
+        pos, inten = seed_peaks(
+            self._volume,
+            blur_sigmas=blur_sigmas,
+            threshold=threshold,
+            threshold_fraction=threshold_fraction,
+            min_distance=min_distance,
+            max_peaks=max_peaks,
+            progress=progress,
+        )
+        self._positions = pos.to(device=self.device, dtype=self._dtype)
+        self._raw_intensity = _inverse_softplus(inten.clamp(min=1e-6)).to(
+            device=self.device, dtype=self._dtype
+        )
+        self._raw_sigma = _inverse_softplus(
+            torch.full((pos.shape[0],), self._sigma_init, dtype=self._dtype, device=self.device)
+        )
+        return self
+
+    # ------------------------------------------------------------------ #
+    # Refinement (joint gradient optimization + site add / remove / merge)
+    # ------------------------------------------------------------------ #
+    def refine(
+        self,
+        num_iterations: int = 100,
+        learning_rate: float = 0.05,
+        loss: str = "huber",
+        sigma_bounds: tuple[float, float] | None = None,
+        intensity_min: float | None = None,
+        add: bool = True,
+        remove: bool = True,
+        merge: bool = True,
+        add_threshold_fraction: float = 0.98,
+        min_neighbors: int = 2,
+        isolation_radius: float | None = None,
+        update_every: int = 25,
+        min_distance: float | None = None,
+        progress: bool = True,
+    ) -> "Atoms":
+        """Refine the atomic model by joint gradient descent.
+
+        All site parameters (positions, intensities, widths) and the background
+        are optimized together against the measured volume.  By default the site
+        set is also maintained during refinement: new sites are *added* at peaks
+        of the residual, weak sites are *removed*, and duplicates are *merged*.
+
+        Parameters
+        ----------
+        num_iterations : int
+            Number of Adam steps.
+        learning_rate : float
+            Base step size (positions, in voxels).  Intensity/background steps are
+            scaled internally by the data magnitude; widths use half this rate.
+        loss : {"huber", "mse"}
+            Data-fidelity term.  Huber is robust to reconstruction artifacts.
+        sigma_bounds : tuple[float, float] or None
+            Hard (min, max) bounds on widths in voxels.  Default
+            ``(0.5, 2.0) * sigma_init``.
+        intensity_min : float or None
+            Sites dimmer than this are removed.  Default 10% of the median site
+            intensity.
+        add, remove, merge : bool
+            Enable adding sites at residual peaks / removing weak (low-intensity)
+            sites / merging close sites during refinement.
+        add_threshold_fraction : float, default 0.98
+            Detection quantile (0-1) for ``add``: lower values recover weaker atoms
+            from the residual (raise toward 1 to add only obvious ones).  Pair with
+            ``min_neighbors`` to reject the extra noise this admits.
+        min_neighbors : int, default 2
+            Remove sites with fewer than this many neighbors within
+            ``isolation_radius`` -- atoms do not float alone, so isolated detections
+            are almost always false positives.  This is what makes a low
+            ``add_threshold_fraction`` usable for finding weak atoms (real weak
+            atoms sit on the lattice and are kept; isolated noise is dropped).  Set
+            0 to disable.
+        isolation_radius : float or None
+            Neighbor-search radius (voxels) for ``min_neighbors``.  Default
+            ``1.5 *`` the estimated nearest-neighbor spacing.
+        update_every : int
+            Apply the add/remove/merge maintenance every this many iterations.
+        min_distance : float or None
+            Minimum site spacing for merge/add, in voxels.  Default
+            ``2 * sigma_init``.
+        progress : bool
+            Show a tqdm progress bar (loss and site count).
+        """
+        if self.num_sites == 0:
+            raise RuntimeError("No sites to refine; call find_initial() first.")
+        if loss not in ("huber", "mse"):
+            raise ValueError("loss must be 'huber' or 'mse'.")
+        sig_lo, sig_hi = sigma_bounds or (0.5 * self._sigma_init, 2.0 * self._sigma_init)
+        if min_distance is None:
+            try:
+                min_distance = 0.75 * self.estimate_spacing()
+            except RuntimeError:
+                min_distance = 2.0 * self._sigma_init
+        self._intensity_scale = float(self._intensity.detach().median().clamp(min=1e-3))
+        if intensity_min is None:
+            intensity_min = 0.1 * self._intensity_scale
+        delta = self._intensity_scale
+
+        self._enable_grad()
+        opt = self._build_optimizer(learning_rate)
+        bar = tqdm(range(num_iterations), desc="refine", disable=not progress)
+        for it in bar:
+            opt.zero_grad(set_to_none=True)
+            model = self.render()
+            if loss == "huber":
+                data_term = F.huber_loss(model, self._volume, delta=delta)
+            else:
+                data_term = F.mse_loss(model, self._volume)
+            data_term.backward()
+            opt.step()
+            # Hard constraints: widths in band, positions inside the volume.
+            with torch.no_grad():
+                clamped = F.softplus(self._raw_sigma).clamp(sig_lo, sig_hi)
+                self._raw_sigma.copy_(_inverse_softplus(clamped))
+                self._positions.clamp_(min=0.0)
+                for a in range(3):
+                    self._positions[:, a].clamp_(max=float(self._shape[a] - 1))
+            # Periodic site-set maintenance.
+            if update_every and (it + 1) % update_every == 0 and it + 1 < num_iterations:
+                changed = False
+                with torch.no_grad():
+                    if add:
+                        changed |= self.add_sites(
+                            min_distance=min_distance, threshold_fraction=add_threshold_fraction
+                        )
+                    if merge:
+                        changed |= self.merge_sites(min_distance)
+                    if min_neighbors > 0:
+                        changed |= self.remove_isolated(isolation_radius, min_neighbors)
+                    if remove:
+                        changed |= self.remove_sites(intensity_min, sigma_bounds=(sig_lo, sig_hi))
+                if changed:
+                    self._enable_grad()
+                    opt = self._build_optimizer(learning_rate)
+            bar.set_postfix(loss=f"{float(data_term.detach()):.4g}", n=self.num_sites)
+        self._disable_grad()
+        return self
+
+    def _enable_grad(self) -> None:
+        for p in (self._positions, self._raw_intensity, self._raw_sigma):
+            p.requires_grad_(True)
+        if self._has_background:
+            self._raw_background.requires_grad_(True)
+
+    def _disable_grad(self) -> None:
+        for p in (self._positions, self._raw_intensity, self._raw_sigma):
+            p.requires_grad_(False)
+        if self._has_background:
+            self._raw_background.requires_grad_(False)
+
+    def _build_optimizer(self, learning_rate: float) -> torch.optim.Optimizer:
+        # Per-group rates: positions in voxels, intensity/background scaled by the
+        # data magnitude, widths slower for stability.
+        scale = getattr(self, "_intensity_scale", 1.0)
+        groups = [
+            {"params": [self._positions], "lr": learning_rate},
+            {"params": [self._raw_sigma], "lr": 0.5 * learning_rate},
+            {"params": [self._raw_intensity], "lr": learning_rate * scale},
+        ]
+        if self._has_background:
+            groups.append({"params": [self._raw_background], "lr": learning_rate * scale})
+        return torch.optim.Adam(groups)
+
+    def remove_sites(self, intensity_min: float, sigma_bounds=None) -> bool:
+        """Remove sites dimmer than ``intensity_min`` (and outside ``sigma_bounds``).
+
+        Returns True if any site was removed.
+        """
+        keep = self._intensity.detach() >= intensity_min
+        if sigma_bounds is not None:
+            sig = self._sigma.detach()
+            keep = keep & (sig >= sigma_bounds[0]) & (sig <= sigma_bounds[1])
+        if bool(keep.all()):
+            return False
+        self._positions = self._positions.detach()[keep]
+        self._raw_intensity = self._raw_intensity.detach()[keep]
+        self._raw_sigma = self._raw_sigma.detach()[keep]
+        return True
+
+    def merge_sites(self, min_distance: float) -> bool:
+        """Merge sites closer than ``min_distance`` voxels, keeping the brightest.
+
+        Returns True if any site was merged away.
+        """
+        if self.num_sites < 2:
+            return False
+        from scipy.spatial import cKDTree
+
+        pos = self._positions.detach().cpu().numpy()
+        inten = self._intensity.detach().cpu().numpy()
+        neighbors = cKDTree(pos).query_ball_point(pos, r=float(min_distance))
+        suppressed = np.zeros(pos.shape[0], dtype=bool)
+        for a in np.argsort(inten)[::-1]:
+            if suppressed[a]:
+                continue
+            for b in neighbors[a]:
+                if b != a and inten[b] <= inten[a]:
+                    suppressed[b] = True
+        if not suppressed.any():
+            return False
+        keep = torch.as_tensor(np.nonzero(~suppressed)[0], device=self.device)
+        self._positions = self._positions.detach()[keep]
+        self._raw_intensity = self._raw_intensity.detach()[keep]
+        self._raw_sigma = self._raw_sigma.detach()[keep]
+        return True
+
+    def remove_isolated(self, radius: float | None = None, min_neighbors: int = 3) -> bool:
+        """Remove isolated sites (fewer than ``min_neighbors`` neighbors within ``radius``).
+
+        Counts how many other sites lie within ``radius`` voxels of each site and
+        drops those below ``min_neighbors``.  Physically, atoms do not float alone
+        in vacuum, so isolated detections are almost always false positives.  This
+        is also what makes detecting *weak* atoms practical: lower the
+        ``find_initial`` / ``add`` threshold to admit weak sites (and noise), then
+        remove the noise here -- real weak atoms sit on the lattice with many
+        neighbors and are kept, while spurious peaks are isolated and removed.
+
+        Parameters
+        ----------
+        radius : float or None
+            Neighbor-search radius in voxels.  Default ``1.5 *`` the estimated
+            nearest-neighbor spacing (:meth:`estimate_spacing`).
+        min_neighbors : int, default 3
+            Minimum neighbors within ``radius`` required to keep a site.
+
+        Returns
+        -------
+        bool
+            True if any site was removed.
+        """
+        if self.num_sites < 2:
+            return False
+        if radius is None:
+            radius = 1.5 * self.estimate_spacing()
+        from scipy.spatial import cKDTree
+
+        pts = self._positions.detach().cpu().numpy()
+        counts = cKDTree(pts).query_ball_point(pts, r=float(radius), return_length=True)
+        keep = (counts - 1) >= min_neighbors  # subtract the site's own match
+        if bool(keep.all()):
+            return False
+        keep_t = torch.as_tensor(np.nonzero(keep)[0], device=self.device)
+        self._positions = self._positions.detach()[keep_t]
+        self._raw_intensity = self._raw_intensity.detach()[keep_t]
+        self._raw_sigma = self._raw_sigma.detach()[keep_t]
+        return True
+
+    def add_sites(
+        self, min_distance: float, threshold_fraction: float = 0.999, max_add: int | None = None
+    ) -> bool:
+        """Add new sites at peaks of the positive residual (densification).
+
+        Detects peaks in ``volume - model`` (clamped to >= 0) that lie at least
+        ``min_distance`` voxels from existing sites, and appends them.  Used during
+        :meth:`refine` to recover atoms the current model is missing.
+
+        Parameters
+        ----------
+        min_distance : float
+            Minimum spacing in voxels, both among new sites and from existing ones.
+        threshold_fraction : float, default 0.999
+            Detection quantile (0-1) on the residual response (see
+            :func:`seed_peaks`); high by default so only clear, missed atoms are
+            added.
+        max_add : int or None
+            If set, cap the number of sites added in this call.
+
+        Returns
+        -------
+        bool
+            True if any site was added.
+        """
+        residual = (self._volume - self.render()).detach().clamp(min=0.0)
+        blur = (0.75 * self._sigma_init, 1.5 * self._sigma_init)
+        new_pos, new_int = seed_peaks(
+            residual,
+            blur_sigmas=blur,
+            threshold_fraction=threshold_fraction,
+            min_distance=min_distance,
+            progress=False,
+        )
+        if new_pos.shape[0] == 0:
+            return False
+        if self.num_sites > 0:
+            from scipy.spatial import cKDTree
+
+            existing = self._positions.detach().cpu().numpy()
+            dist, _ = cKDTree(existing).query(new_pos.cpu().numpy(), k=1)
+            far = torch.as_tensor(dist >= float(min_distance))
+            new_pos, new_int = new_pos[far], new_int[far]
+        if new_pos.shape[0] == 0:
+            return False
+        if max_add is not None:
+            new_pos, new_int = new_pos[:max_add], new_int[:max_add]
+        new_pos = new_pos.to(device=self.device, dtype=self._dtype)
+        new_raw_int = _inverse_softplus(
+            new_int.to(device=self.device, dtype=self._dtype).clamp(min=1e-6)
+        )
+        new_raw_sig = _inverse_softplus(
+            torch.full((new_pos.shape[0],), self._sigma_init, dtype=self._dtype, device=self.device)
+        )
+        self._positions = torch.cat([self._positions.detach(), new_pos], dim=0)
+        self._raw_intensity = torch.cat([self._raw_intensity.detach(), new_raw_int], dim=0)
+        self._raw_sigma = torch.cat([self._raw_sigma.detach(), new_raw_sig], dim=0)
+        return True
+
+    # ------------------------------------------------------------------ #
+    # Rendering / decomposition
+    # ------------------------------------------------------------------ #
+    def _render_atoms(self) -> Tensor:
+        return render_isotropic(
+            self._positions, self._intensity, self._sigma, self._shape, self._window_radius()
+        )
+
+    def _render_background(self) -> Tensor:
+        if not self._has_background:
+            return torch.zeros(self._shape, dtype=self._dtype, device=self.device)
+        return render_background(self._background_control, self._bases)
+
+    def render(self) -> Tensor:
+        """Render the full model ``volume_atoms + volume_background`` as a tensor."""
+        return self._render_atoms() + self._render_background()
+
+    def _as_dataset(self, tensor: Tensor, name: str) -> Dataset3d:
+        arr = tensor.detach().to("cpu").numpy().astype(np.float32)
+        return Dataset3d.from_array(
+            arr,
+            name=name,
+            origin=self._origin,
+            sampling=self._sampling,
+            units=self._units,
+            signal_units=self._signal_units,
+        )
+
+    @property
+    def volume(self) -> Dataset3d:
+        """The source volume."""
+        return self._source
+
+    @property
+    def volume_atoms(self) -> Dataset3d:
+        """Rendered atomic part (>= 0)."""
+        return self._as_dataset(self._render_atoms(), "volume_atoms")
+
+    @property
+    def volume_background(self) -> Dataset3d:
+        """Rendered smooth background (>= 0)."""
+        return self._as_dataset(self._render_background(), "volume_background")
+
+    @property
+    def volume_model(self) -> Dataset3d:
+        """Full model ``volume_atoms + volume_background``."""
+        return self._as_dataset(self.render(), "volume_model")
+
+    @property
+    def residual(self) -> Dataset3d:
+        """Source volume minus the model."""
+        return self._as_dataset(self._volume - self.render(), "residual")
+
+    # ------------------------------------------------------------------ #
+    # Sites as a Vector (physical units)
+    # ------------------------------------------------------------------ #
+    @property
+    def sites(self) -> Vector:
+        """Atomic sites as a :class:`Vector` with fields x, y, z, intensity, sigma.
+
+        Positions and sigma are in the source volume's physical units (sigma uses
+        the mean sampling).
+        """
+        fields = ["x", "y", "z", "intensity", "sigma"]
+        units = [self._units[0], self._units[1], self._units[2], self._signal_units, self._units[0]]
+        v = Vector.from_shape(shape=(), fields=fields, units=units, name="atoms")
+        n = self.num_sites
+        if n == 0:
+            return v
+        pos_vox = self._positions.detach().to("cpu").numpy()
+        pos_phys = self._origin[None, :] + pos_vox * self._sampling[None, :]
+        inten = self._intensity.detach().to("cpu").numpy()
+        sigma_phys = self._sigma.detach().to("cpu").numpy() * float(self._sampling.mean())
+        data = np.column_stack([pos_phys, inten, sigma_phys]).astype(np.float32)
+        v[...] = data
+        return v
+
+    # ------------------------------------------------------------------ #
+    # Interactive 3D widget
+    # ------------------------------------------------------------------ #
+    def show_3d_atoms(self, **kwargs):
+        """Open the interactive 3D slice + atom-overlay widget.
+
+        Renders an orthogonal slice (xy / xz / yz) through the volume with the
+        atomic sites overlaid (marker size scaled by intensity, opacity fading
+        with distance from the slice).  Sites are passed in voxel coordinates so
+        they register with the volume.  Requires the ``quantem.widget`` package;
+        keyword arguments are forwarded to :class:`quantem.widget.Show3DAtoms`.
+        """
+        from quantem.widget import Show3DAtoms
+
+        if self.num_sites:
+            sites = np.column_stack(
+                [
+                    self._positions.detach().to("cpu").numpy(),
+                    self._intensity.detach().to("cpu").numpy(),
+                    self._sigma.detach().to("cpu").numpy(),
+                ]
+            ).astype(np.float32)
+        else:
+            sites = np.zeros((0, 5), dtype=np.float32)
+        volume = self._volume.detach().to("cpu").numpy().astype(np.float32)
+        kwargs.setdefault("title", str(self._source.name))
+        kwargs.setdefault("sampling", tuple(float(s) for s in self._sampling))
+        return Show3DAtoms(volume, sites=sites, **kwargs)
+
+    def __repr__(self) -> str:
+        return (
+            f"quantem.Atoms(model={self._model}, num_sites={self.num_sites}, "
+            f"volume_shape={self._shape}, device={self.device}, "
+            f"background={self._has_background})"
+        )
diff --git a/src/quantem/tomography/dataset_models.py b/src/quantem/tomography/dataset_models.py
index 2255636c0..abc3ef4c7 100644
--- a/src/quantem/tomography/dataset_models.py
+++ b/src/quantem/tomography/dataset_models.py
@@ -1,3 +1,4 @@
+import math
 from abc import abstractmethod
 from dataclasses import dataclass
 from typing import Any
@@ -10,7 +11,7 @@
 from quantem.core.datastructures.dataset3d import Dataset3d
 from quantem.core.io.serialize import AutoSerialize
 from quantem.core.ml.constraints import BaseConstraints, Constraints
-from quantem.core.ml.optimizer_mixin import OptimizerMixin
+from quantem.core.ml.optimizer_mixin import OptimizerMixin, OptimizerParams, OptimizerParamsType
 from quantem.tomography.utils import tv_loss_1d
 
 # --- Constraints ---
@@ -152,6 +153,158 @@ class DatasetValue:
     )
 
 
+@dataclass(frozen=True)
+class PixelHoldoutSplit:
+    """Flat pixel indices for train and held-out validation rays."""
+
+    train_indices: torch.Tensor
+    val_indices: torch.Tensor
+    val_fg_indices: torch.Tensor
+    val_bg_indices: torch.Tensor
+
+
+def _allocate_pixel_holdout_counts(group_sizes: list[int], n_holdout: int) -> list[int]:
+    """Allocate an exact holdout count across groups by largest remainder."""
+    if n_holdout < 0:
+        raise ValueError("n_holdout must be >= 0.")
+    total = sum(group_sizes)
+    if n_holdout > total:
+        raise ValueError("n_holdout cannot exceed the total number of pixels.")
+    if total == 0 or n_holdout == 0:
+        return [0 for _ in group_sizes]
+
+    quotas = [n_holdout * size / total for size in group_sizes]
+    counts = [min(size, int(quota)) for size, quota in zip(group_sizes, quotas)]
+    remaining = n_holdout - sum(counts)
+    order = sorted(
+        range(len(group_sizes)),
+        key=lambda i: (quotas[i] - int(quotas[i]), group_sizes[i]),
+        reverse=True,
+    )
+    while remaining > 0:
+        progressed = False
+        for i in order:
+            if counts[i] < group_sizes[i]:
+                counts[i] += 1
+                remaining -= 1
+                progressed = True
+                if remaining == 0:
+                    break
+        if not progressed:
+            raise RuntimeError("Could not allocate holdout counts.")
+    return counts
+
+
+def _sample_intensity_strata(
+    *,
+    indices: torch.Tensor,
+    intensities: torch.Tensor,
+    n_holdout: int,
+    seed: int,
+    num_strata: int,
+) -> torch.Tensor:
+    """Sample held-out indices from intensity-quantile strata."""
+    if n_holdout == 0 or indices.numel() == 0:
+        return torch.empty(0, dtype=torch.long)
+
+    if n_holdout > indices.numel():
+        raise ValueError("n_holdout cannot exceed the number of candidate pixels.")
+    n_strata = max(1, min(int(num_strata), int(indices.numel())))
+    order = torch.argsort(intensities[indices], stable=True)
+    sorted_indices = indices[order]
+    strata = [s for s in torch.tensor_split(sorted_indices, n_strata) if s.numel() > 0]
+    counts = _allocate_pixel_holdout_counts([int(s.numel()) for s in strata], n_holdout)
+
+    generator = torch.Generator(device="cpu")
+    generator.manual_seed(int(seed))
+    selected: list[torch.Tensor] = []
+    for stratum, count in zip(strata, counts):
+        if count == 0:
+            continue
+        perm = torch.randperm(stratum.numel(), generator=generator)
+        selected.append(stratum[perm[:count]])
+
+    if not selected:
+        return torch.empty(0, dtype=torch.long)
+    return torch.cat(selected).to(dtype=torch.long)
+
+
+def build_pixel_holdout_split(
+    tilt_stack: Dataset3d | NDArray | torch.Tensor,
+    holdout_fraction: float,
+    holdout_seed: int = 0,
+    *,
+    num_strata: int = 8,
+    foreground_threshold: float = 0.0,
+) -> PixelHoldoutSplit:
+    """Build a seeded foreground-aware train/holdout split over flat pixel indices.
+
+    The split is rank-independent: all work happens on detached CPU tensors with a
+    local generator seeded only by ``holdout_seed``. Foreground/background groups are
+    split proportionally, then each group is sampled from intensity-quantile strata.
+    """
+    if not 0.0 <= holdout_fraction < 1.0:
+        raise ValueError("holdout_fraction must satisfy 0 <= holdout_fraction < 1.")
+    stack = torch.as_tensor(tilt_stack).detach().cpu()
+    intensities = stack.to(torch.float32).abs().flatten()
+    n_pixels = int(intensities.numel())
+    all_indices = torch.arange(n_pixels, dtype=torch.long)
+    n_holdout = int(n_pixels * float(holdout_fraction))
+
+    if n_holdout == 0:
+        empty = torch.empty(0, dtype=torch.long)
+        return PixelHoldoutSplit(
+            train_indices=all_indices,
+            val_indices=empty,
+            val_fg_indices=empty,
+            val_bg_indices=empty,
+        )
+
+    fg_mask = intensities > float(foreground_threshold)
+    fg_indices = all_indices[fg_mask]
+    bg_indices = all_indices[~fg_mask]
+    bg_count, fg_count = _allocate_pixel_holdout_counts(
+        [int(bg_indices.numel()), int(fg_indices.numel())], n_holdout
+    )
+
+    if (
+        n_holdout >= 2
+        and bg_indices.numel() > 0
+        and fg_indices.numel() > 0
+        and (bg_count == 0 or fg_count == 0)
+    ):
+        if bg_count == 0 and fg_count > 1:
+            bg_count, fg_count = 1, fg_count - 1
+        elif fg_count == 0 and bg_count > 1:
+            bg_count, fg_count = bg_count - 1, 1
+
+    bg_val = _sample_intensity_strata(
+        indices=bg_indices,
+        intensities=intensities,
+        n_holdout=bg_count,
+        seed=int(holdout_seed) * 2 + 1,
+        num_strata=num_strata,
+    )
+    fg_val = _sample_intensity_strata(
+        indices=fg_indices,
+        intensities=intensities,
+        n_holdout=fg_count,
+        seed=int(holdout_seed) * 2 + 2,
+        num_strata=num_strata,
+    )
+    val_indices = torch.sort(torch.cat([bg_val, fg_val])).values
+    train_mask = torch.ones(n_pixels, dtype=torch.bool)
+    train_mask[val_indices] = False
+    train_indices = all_indices[train_mask]
+
+    return PixelHoldoutSplit(
+        train_indices=train_indices,
+        val_indices=val_indices,
+        val_fg_indices=torch.sort(fg_val).values,
+        val_bg_indices=torch.sort(bg_val).values,
+    )
+
+
 class TomographyDatasetBase(AutoSerialize, OptimizerMixin, nn.Module):
     """
     Base tomography dataset class for all tomography datasets to inherit from.
@@ -177,9 +330,7 @@ def __init__(
         if _token is not self._token:
             raise RuntimeError("Use TomographyPixDataset.from_* to instantiate this class.")
 
-        if not (
-            tilt_stack.shape[0] < tilt_stack.shape[1] or tilt_stack.shape[0] < tilt_stack.shape[2]
-        ):
+        if tilt_stack.shape[0] != len(tilt_angles):
             raise ValueError(
                 "The number of tilt projections should be in the first dimension of the dataset."
             )
@@ -189,6 +340,12 @@ def __init__(
         if type(tilt_angles) is not torch.Tensor:
             tilt_angles = torch.from_numpy(tilt_angles)
         max_val = torch.quantile(tilt_stack, 0.95)
+        # A sparse stack (>95% zeros) has a zero 95th quantile; dividing by it
+        # would turn the targets into inf/NaN and poison the first backward.
+        if max_val <= 0:
+            max_val = tilt_stack.abs().max()
+        if max_val <= 0:
+            raise ValueError("tilt_stack is all zeros; cannot normalize.")
 
         # Tilt stack normalization
         tilt_stack = tilt_stack / max_val
@@ -240,6 +397,21 @@ def get_optimization_parameters(self) -> dict[str, list[torch.Tensor]]:
         """
         return {self.DEFAULT_OPTIMIZER_KEY: list(self.parameters())}
 
+    def _materialize_pose_parameters(self, device: str | torch.device):
+        """Create the learnable pose parameters, or move the existing ones.
+
+        Once the parameters exist, their *current* (possibly trained) values are
+        moved; the initial-value buffers are only used on first materialization.
+        Rebuilding from the buffers on every call silently reset learned poses
+        whenever the dataset changed device (e.g. ``from_file(...).to(device)``).
+        """
+        z1 = self._z1_params.data if hasattr(self, "_z1_params") else self._z1_angles
+        z3 = self._z3_params.data if hasattr(self, "_z3_params") else self._z3_angles
+        shifts = self._shifts_params.data if hasattr(self, "_shifts_params") else self._shifts
+        self._z1_params = nn.Parameter(z1.detach().to(device))
+        self._z3_params = nn.Parameter(z3.detach().to(device))
+        self._shifts_params = nn.Parameter(shifts.detach().to(device))
+
     # --- Forward pass ---
     @abstractmethod
     def forward(
@@ -302,12 +474,11 @@ def reference_tilt_idx(self, reference_tilt_idx: int):
 
     @property
     def learnable_tilts(self) -> int:
+        # Derived from the tilt series (all tilts minus the fixed reference); there is
+        # deliberately no setter -- the old one wrote a private attribute this getter
+        # never read, so assignments appeared to succeed while doing nothing.
         return self.tilt_angles.shape[0] - 1
 
-    @learnable_tilts.setter
-    def learnable_tilts(self, learnable_tilts: int):
-        self._learnable_tilts = learnable_tilts
-
     @property
     def z1_params(self) -> torch.nn.Parameter:
         return self._z1_params
@@ -362,7 +533,7 @@ def __init__(self, *args, **kwargs):
         )
 
     def apply_soft_constraints(self) -> torch.Tensor:
-        soft_loss = torch.tensor(0.0, device=self.z1_params.device)
+        soft_loss = torch.zeros((), device=self.z1_params.device)
         if self.constraints.tv_zs > 0:
             tv_loss_zs = tv_loss_1d(self.z1_params)
             tv_loss_zs += tv_loss_1d(self.z3_params)
@@ -429,9 +600,7 @@ def to(self, device: str | torch.device):
         self.tilt_stack = self.tilt_stack.to(device)
         self.tilt_angles = self.tilt_angles.to(device)
 
-        self._z1_params = nn.Parameter(self._z1_angles.to(device))
-        self._z3_params = nn.Parameter(self._z3_angles.to(device))
-        self._shifts_params = nn.Parameter(self._shifts.to(device))
+        self._materialize_pose_parameters(device)
 
         self._z1_ref = self._z1_ref.to(device)
         self._z3_ref = self._z3_ref.to(device)
@@ -440,6 +609,114 @@ def to(self, device: str | torch.device):
         self.device = device
 
 
+class DeviceBatchSampler:
+    """Epoch iterator that builds INR training batches directly on a device.
+
+    Replaces the per-pixel DataLoader path: the tilt stack and angles are
+    made resident on ``device`` once, so producing a batch is index
+    arithmetic plus two tensor lookups instead of ``batch_size`` Python
+    ``__getitem__`` calls, a collate, and a host-to-device copy per step.
+    On a GPU this removes the CPU dataloader bottleneck entirely.
+
+    Yields the same batch dicts as ``TomographyINRDataset.__getitem__``
+    under a DataLoader collate (``projection_idx``, ``pixel_i``,
+    ``pixel_j``, ``phi``, ``target_value``), with the train loader's
+    ``drop_last=True`` semantics.
+
+    Distributed runs: pass ``rank``/``world_size`` and every rank derives
+    the *same* epoch permutation from ``seed + epoch`` (CPU generator, so
+    it is identical across ranks and reproducible), then takes an
+    equal-size contiguous shard — equal so per-rank batch counts match and
+    DDP gradient sync cannot hang on a ragged tail. The training loop's
+    ``sampler.set_epoch(epoch)`` drives reshuffling, exactly like
+    ``DistributedSampler``; without ``set_epoch`` the epoch advances
+    automatically on each ``__iter__``.
+    """
+
+    def __init__(
+        self,
+        dset: "TomographyINRDataset",
+        batch_size: int,
+        device: torch.device | str,
+        indices: torch.Tensor | None = None,
+        shuffle: bool = True,
+        rank: int = 0,
+        world_size: int = 1,
+        seed: int = 0,
+        drop_last: bool = True,
+    ):
+        self.batch_size = batch_size
+        self.device = torch.device(device)
+        self.shuffle = shuffle
+        self.rank = rank
+        self.world_size = world_size
+        self.seed = seed
+        self.drop_last = drop_last
+        self._epoch = 0
+        self._stack = dset.tilt_stack.to(self.device)
+        self._angles = dset.tilt_angles.to(self.device)
+        self._angles_per_row = (
+            dset.tilt_angles_per_row.to(self.device)
+            if dset.tilt_angles_per_row is not None
+            else None
+        )
+        self._angles_per_col = (
+            dset.tilt_angles_per_col.to(self.device)
+            if dset.tilt_angles_per_col is not None
+            else None
+        )
+        self._s1 = dset.tilt_stack.shape[1]
+        self._s2 = dset.tilt_stack.shape[2]
+        if indices is None:
+            indices = torch.arange(len(dset), dtype=torch.int64)
+        self._indices = indices.to(self.device)
+        self._per_rank = len(self._indices) // world_size
+
+    def set_epoch(self, epoch: int) -> None:
+        """Set the epoch used to seed this epoch's shared permutation."""
+        self._epoch = epoch
+
+    def __len__(self) -> int:
+        if self.drop_last:
+            return self._per_rank // self.batch_size
+        return (self._per_rank + self.batch_size - 1) // self.batch_size
+
+    def _epoch_shard(self) -> torch.Tensor:
+        idx = self._indices
+        if self.shuffle:
+            g = torch.Generator()
+            g.manual_seed(self.seed + self._epoch)
+            perm = torch.randperm(len(idx), generator=g).to(self.device)
+            idx = idx[perm]
+            self._epoch += 1  # auto-advance; set_epoch overrides per epoch
+        if self.world_size > 1:
+            idx = idx[self.rank * self._per_rank : (self.rank + 1) * self._per_rank]
+        return idx
+
+    def __iter__(self):
+        idx = self._epoch_shard()
+        per_proj = self._s1 * self._s2
+        for k in range(len(self)):
+            sel = idx[k * self.batch_size : min((k + 1) * self.batch_size, len(idx))]
+            proj = sel // per_proj
+            rem = sel - proj * per_proj
+            pixel_i = rem // self._s2
+            pixel_j = rem - pixel_i * self._s2
+            if self._angles_per_row is not None:
+                phi = self._angles_per_row[proj, pixel_i]
+            elif self._angles_per_col is not None:
+                phi = self._angles_per_col[proj, pixel_j]
+            else:
+                phi = self._angles[proj]
+            yield {
+                "projection_idx": proj,
+                "pixel_i": pixel_i,
+                "pixel_j": pixel_j,
+                "phi": phi,
+                "target_value": self._stack[proj, pixel_i, pixel_j],
+            }
+
+
 class TomographyINRDataset(TomographyDatasetConstraints, Dataset):
     """
     Dataset class for INR-based tomography.
@@ -454,14 +731,118 @@ def __init__(
         self,
         tilt_stack: Dataset3d | NDArray | torch.Tensor,
         tilt_angles: NDArray | torch.Tensor,
+        tilt_angles_per_row: NDArray | torch.Tensor | None = None,
         learn_shift: bool = True,
         learn_tilt_axis: bool = True,
+        ray_sampling: str = "box_fixed_ds",
+        ray_ds: float | None = None,
         seed: int = 42,
+        tilt_angles_per_col: NDArray | torch.Tensor | None = None,
         _token: object | None = None,
     ):
         super().__init__(tilt_stack, tilt_angles, learn_shift, learn_tilt_axis, _token=_token)
 
+        if tilt_angles_per_row is not None and tilt_angles_per_col is not None:
+            raise ValueError("tilt_angles_per_row and tilt_angles_per_col are mutually exclusive")
+        if tilt_angles_per_row is not None:
+            if type(tilt_angles_per_row) is not torch.Tensor:
+                tilt_angles_per_row = torch.from_numpy(tilt_angles_per_row)
+            if tilt_angles_per_row.ndim != 2 or tuple(tilt_angles_per_row.shape) != tuple(
+                tilt_stack.shape[:2]
+            ):
+                raise ValueError(
+                    "tilt_angles_per_row must have shape "
+                    f"{tuple(tilt_stack.shape[:2])}, got {tuple(tilt_angles_per_row.shape)}."
+                )
+        self.tilt_angles_per_row = tilt_angles_per_row
+        if tilt_angles_per_col is not None:
+            if type(tilt_angles_per_col) is not torch.Tensor:
+                tilt_angles_per_col = torch.from_numpy(tilt_angles_per_col)
+            expected_shape = (tilt_stack.shape[0], tilt_stack.shape[2])
+            if tilt_angles_per_col.ndim != 2 or tuple(tilt_angles_per_col.shape) != tuple(
+                expected_shape
+            ):
+                raise ValueError(
+                    "tilt_angles_per_col must have shape "
+                    f"{tuple(expected_shape)}, got {tuple(tilt_angles_per_col.shape)}."
+                )
+        self.tilt_angles_per_col = tilt_angles_per_col
+
+        # --- Ray-sampling scheme ---
+        # "legacy"       : detector-frame z in [-1, 1] swept with a FIXED count of points,
+        #                  then rotated (create_batch_rays / transform_batch_rays). The
+        #                  segment is rotated with the object, so long diagonal chords are
+        #                  clipped and other rays waste samples outside the [-1,1]^3 cube.
+        # "box_fixed_ds" : per-ray ray-box intersection with the [-1,1]^3 cube, then a
+        #                  CONSTANT physical step `ds` along the true chord (variable count
+        #                  per ray). Gives identical physical sampling spacing on every ray
+        #                  regardless of tilt -- the consistent discretization of the line
+        #                  integral. Ragged, so integrate_rays uses a scatter-add.
+        self.ray_sampling: str = ray_sampling
+        # Physical step between samples for "box_fixed_ds". None => derive per call from
+        # num_samples_per_ray as 2/(num_samples_per_ray - 1) so the existing samples_per_ray
+        # knob stays meaningful (a centered ray reproduces the legacy sample count).
+        self.ray_ds: float | None = ray_ds
+        # Per-batch ragged metadata stashed by get_coords for integrate_rays to consume.
+        self._ray_meta: dict[str, torch.Tensor] | None = None
+
+    @classmethod
+    def from_data(
+        cls,
+        tilt_stack: Dataset3d | NDArray | torch.Tensor,
+        tilt_angles: NDArray | torch.Tensor,
+        tilt_angles_per_row: NDArray | torch.Tensor | None = None,
+        learn_shift: bool = True,
+        learn_tilt_axis: bool = True,
+        ray_sampling: str = "box_fixed_ds",
+        ray_ds: float | None = None,
+        tilt_angles_per_col: NDArray | torch.Tensor | None = None,
+    ):
+
+        if ray_sampling == "box_fixed_ds":
+            if ray_ds is None:
+                ray_ds = 2.0 / max(tilt_stack.shape)
+            else:
+                ray_ds = float(ray_ds)
+
+        return cls(
+            tilt_stack=tilt_stack,
+            tilt_angles=tilt_angles,
+            tilt_angles_per_row=tilt_angles_per_row,
+            tilt_angles_per_col=tilt_angles_per_col,
+            learn_shift=learn_shift,
+            learn_tilt_axis=learn_tilt_axis,
+            ray_sampling=ray_sampling,
+            ray_ds=ray_ds,
+            _token=cls._token,
+        )
+
     # --- Forward Pass w/ Params Method for OptimizerMixin ---
+    def get_optimization_parameters(self) -> dict[str, list[torch.Tensor]]:
+        """Return independently tunable shift and tilt-axis parameter groups."""
+        groups = {}
+        if self.learn_shift:
+            groups["pose_shift"] = [self._shifts_params]
+        if self.learn_tilt_axis:
+            groups["pose_tilt_axis"] = [self._z1_params, self._z3_params]
+        return groups
+
+    def _normalize_optimizer_params(
+        self, params: OptimizerParamsType | dict[str, Any]
+    ) -> dict[str, OptimizerParamsType]:
+        """Expand a legacy shared pose optimizer over the active pose groups."""
+        normalized = super()._normalize_optimizer_params(params)
+        if set(normalized) == {self.DEFAULT_OPTIMIZER_KEY}:
+            spec = normalized[self.DEFAULT_OPTIMIZER_KEY]
+            if not isinstance(spec, OptimizerParams.NoneOptimizer):
+                normalized = {
+                    key: spec
+                    for key in ("pose_shift", "pose_tilt_axis")
+                    if (key == "pose_shift" and self.learn_shift)
+                    or (key == "pose_tilt_axis" and self.learn_tilt_axis)
+                }
+        return normalized
+
     def forward(self, dummy_input: Any = None):
         """
         Forward pass for INR-based tomography. In the forward pass, the only parameters that
@@ -497,14 +878,27 @@ def get_coords(
         # target_values = batch["target_value"].to(self.device, non_blocking=True)
         phis = batch["phi"].to(self.device, non_blocking=True)
         projection_indices = batch["projection_idx"].to(self.device, non_blocking=True)
-        with torch.no_grad():
-            batch_ray_coords = self.create_batch_rays(pixel_i, pixel_j, N, num_samples_per_ray)
 
         shifts, z1_params, z3_params = self.forward(None)
         batch_shifts = torch.index_select(shifts, 0, projection_indices)
         batch_z1 = torch.index_select(z1_params, 0, projection_indices)
         batch_z3 = torch.index_select(z3_params, 0, projection_indices)
 
+        if getattr(self, "ray_sampling", "legacy") == "box_fixed_ds":
+            return self._get_coords_box_fixed_ds(
+                pixel_i=pixel_i,
+                pixel_j=pixel_j,
+                phis=phis,
+                batch_z1=batch_z1,
+                batch_z3=batch_z3,
+                batch_shifts=batch_shifts,
+                N=N,
+                num_samples_per_ray=num_samples_per_ray,
+            )
+
+        with torch.no_grad():
+            batch_ray_coords = self.create_batch_rays(pixel_i, pixel_j, N, num_samples_per_ray)
+
         transformed_rays = self.transform_batch_rays(
             batch_ray_coords,
             z1=batch_z1,
@@ -550,41 +944,279 @@ def transform_batch_rays(
         shift_x_norm = (shifts[:, 0:1] * sampling_rate * 2) / (N - 1)
         shift_y_norm = (shifts[:, 1:2] * sampling_rate * 2) / (N - 1)
 
-        rays_x = rays[:, :, 0] - shift_x_norm
-        rays_y = rays[:, :, 1] - shift_y_norm
-        rays_z = rays[:, :, 2]
+        shifted = torch.stack(
+            [rays[:, :, 0] - shift_x_norm, rays[:, :, 1] - shift_y_norm, rays[:, :, 2]],
+            dim=2,
+        )
 
-        theta = torch.deg2rad(-z3).view(-1, 1)
-        cos_t = torch.cos(theta)
-        sin_t = torch.sin(theta)
+        rot = TomographyINRDataset._compose_euler_rotation(z1, x, z3)  # (B, 3, 3)
 
-        rays_x_rot1 = cos_t * rays_x - sin_t * rays_y
-        rays_y_rot1 = sin_t * rays_x + cos_t * rays_y
-        rays_z_rot1 = rays_z
+        return shifted @ rot.transpose(1, 2)
 
-        theta = torch.deg2rad(x).view(-1, 1)
-        cos_t = torch.cos(theta)
-        sin_t = torch.sin(theta)
+    @staticmethod
+    def _compose_euler_rotation(
+        z1: torch.Tensor, x: torch.Tensor, z3: torch.Tensor
+    ) -> torch.Tensor:
+        """Compose the three Euler rotations Rz(-z1) @ Rx(x) @ Rz(-z3) into a single
+        (B, 3, 3) matrix. A point row-vector ``v`` is mapped to the object frame by
+        ``v @ rot.transpose`` (equivalently ``rot @ v``)."""
+        a = torch.deg2rad(-z3).view(-1)
+        b = torch.deg2rad(x).view(-1)
+        g = torch.deg2rad(-z1).view(-1)
+        zero = torch.zeros_like(a)
+        one = torch.ones_like(a)
+
+        cos_a, sin_a = torch.cos(a), torch.sin(a)
+        cos_b, sin_b = torch.cos(b), torch.sin(b)
+        cos_g, sin_g = torch.cos(g), torch.sin(g)
+
+        rot_a = torch.stack(
+            [cos_a, -sin_a, zero, sin_a, cos_a, zero, zero, zero, one], dim=-1
+        ).view(-1, 3, 3)
+        rot_b = torch.stack(
+            [one, zero, zero, zero, cos_b, -sin_b, zero, sin_b, cos_b], dim=-1
+        ).view(-1, 3, 3)
+        rot_g = torch.stack(
+            [cos_g, -sin_g, zero, sin_g, cos_g, zero, zero, zero, one], dim=-1
+        ).view(-1, 3, 3)
+
+        return rot_g @ rot_b @ rot_a  # (B, 3, 3)
+
+    # ------------------------------------------------------------------
+    # Box-intersection, fixed-distance ("box_fixed_ds") ray sampling.
+    # Ported from ray_sampling/geometry.py (ParallelBeamRayProjector), reusing the
+    # dataset's own Euler pose so gradients keep flowing to z1/z3/shifts.
+    # ------------------------------------------------------------------
+    def _resolve_ray_ds(self, num_samples_per_ray: int) -> float:
+        """Physical step `ds` between samples along every ray (object-space units,
+        where the cube spans [-1, 1])."""
+        if self.ray_ds is not None:
+            ds = float(self.ray_ds)
+            if ds <= 0.0:
+                raise ValueError(f"ray_ds must be > 0, got {ds}")
+            return ds
+        # Derive from the samples_per_ray knob so it stays meaningful: an untilted ray
+        # (chord length 2) then gets exactly num_samples_per_ray points, matching legacy.
+        return 2.0 / (num_samples_per_ray - 1)
+
+    def _build_ray_origins_directions(
+        self,
+        *,
+        pixel_i: torch.Tensor,
+        pixel_j: torch.Tensor,
+        z1: torch.Tensor,
+        x: torch.Tensor,
+        z3: torch.Tensor,
+        shifts: torch.Tensor,
+        N: int,
+        sampling_rate: float,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Ray origin/direction in the object frame for each detector pixel.
+
+        Rotates the two legacy endpoints (x, y, +/-1) by the same Euler pose used in
+        transform_batch_rays, so ``point(t) = origin + t * direction`` reproduces the
+        legacy z in [-1, 1] segment at t in [-1, 1] but lets t run past it to capture the
+        full chord through the cube. ``direction`` is unit-norm (rotation of (0,0,1)), so
+        t is true arc length and ``ds`` is a physical distance.
+        """
+        x_coords = (pixel_j / (N - 1)) * 2 - 1
+        y_coords = (pixel_i / (N - 1)) * 2 - 1
+        shift_x_norm = (shifts[:, 0] * sampling_rate * 2) / (N - 1)
+        shift_y_norm = (shifts[:, 1] * sampling_rate * 2) / (N - 1)
+        x_base = x_coords - shift_x_norm
+        y_base = y_coords - shift_y_norm
+
+        rot = self._compose_euler_rotation(z1, x, z3)  # (B, 3, 3)
+        ones = torch.ones_like(x_base)
+        p_plus_local = torch.stack((x_base, y_base, ones), dim=-1)
+        p_minus_local = torch.stack((x_base, y_base, -ones), dim=-1)
+        p_plus = torch.einsum("bij,bj->bi", rot, p_plus_local)
+        p_minus = torch.einsum("bij,bj->bi", rot, p_minus_local)
+
+        origins = 0.5 * (p_plus + p_minus)
+        directions = 0.5 * (p_plus - p_minus)
+        return origins, directions
 
-        rays_x_rot2 = rays_x_rot1
-        rays_y_rot2 = cos_t * rays_y_rot1 - sin_t * rays_z_rot1
-        rays_z_rot2 = sin_t * rays_y_rot1 + cos_t * rays_z_rot1
+    @staticmethod
+    def _compute_ray_box_intersections(
+        origins: torch.Tensor,
+        directions: torch.Tensor,
+        eps: float = 1.0e-8,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Slab method: entry/exit t into the [-1, 1]^3 cube and a validity mask."""
+        parallel = directions.abs() < eps
+        safe_dirs = torch.where(parallel, torch.ones_like(directions), directions)
+
+        t0 = (-1.0 - origins) / safe_dirs
+        t1 = (1.0 - origins) / safe_dirs
+        t_min_dim = torch.minimum(t0, t1)
+        t_max_dim = torch.maximum(t0, t1)
+
+        neg_inf = torch.full_like(t_min_dim, -float("inf"))
+        pos_inf = torch.full_like(t_max_dim, float("inf"))
+        in_slab = (origins >= -1.0) & (origins <= 1.0)
+        valid_parallel = (~parallel) | in_slab
+
+        t_min_dim = torch.where(parallel, neg_inf, t_min_dim)
+        t_max_dim = torch.where(parallel, pos_inf, t_max_dim)
+
+        t_enter = t_min_dim.max(dim=1).values
+        t_exit = t_max_dim.min(dim=1).values
+        valid = valid_parallel.all(dim=1) & (t_exit > t_enter)
+        return t_enter, t_exit, valid
+
+    def _sample_ray_segment_coords_fixed_ds(
+        self,
+        origins: torch.Tensor,
+        directions: torch.Tensor,
+        t_enter: torch.Tensor,
+        t_exit: torch.Tensor,
+        ds: float,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Place points every ``ds`` along each ray's in-cube chord.
+
+        Variable count per ray (padded to the batch max), so a ``sample_mask`` marks the
+        real samples. Returns (coords, sample_mask, n_samples, lengths). The sample COUNT
+        is decided from detached lengths (ceil is non-differentiable anyway); the sample
+        POSITIONS stay differentiable in t_enter/lengths, so pose gradients flow.
+        """
+        lengths = t_exit - t_enter
+        n_samples = torch.ceil(lengths.detach() / float(ds)).to(torch.long) + 1
+        n_samples = torch.clamp(n_samples, min=2)
 
-        theta = torch.deg2rad(-z1).view(-1, 1)
-        cos_t = torch.cos(theta)
-        sin_t = torch.sin(theta)
+        max_samples = int(n_samples.max().item())
+        sample_idx = torch.arange(max_samples, device=origins.device).unsqueeze(0)
+        sample_mask = sample_idx < n_samples.unsqueeze(1)
 
-        rays_x_final = cos_t * rays_x_rot2 - sin_t * rays_y_rot2
-        rays_y_final = sin_t * rays_x_rot2 + cos_t * rays_y_rot2
-        rays_z_final = rays_z_rot2
+        denom = (n_samples - 1).to(origins.dtype).unsqueeze(1)
+        u = sample_idx.to(origins.dtype) / denom
+        t_vals = t_enter.unsqueeze(1) + lengths.unsqueeze(1) * u
+        coords = origins.unsqueeze(1) + t_vals.unsqueeze(2) * directions.unsqueeze(1)
+        return coords, sample_mask, n_samples, lengths
 
-        transformed_rays = torch.stack([rays_x_final, rays_y_final, rays_z_final], dim=2)
+    def _get_coords_box_fixed_ds(
+        self,
+        *,
+        pixel_i: torch.Tensor,
+        pixel_j: torch.Tensor,
+        phis: torch.Tensor,
+        batch_z1: torch.Tensor,
+        batch_z3: torch.Tensor,
+        batch_shifts: torch.Tensor,
+        N: int,
+        num_samples_per_ray: int,
+    ) -> torch.Tensor:
+        origins, directions = self._build_ray_origins_directions(
+            pixel_i=pixel_i,
+            pixel_j=pixel_j,
+            z1=batch_z1,
+            x=phis,
+            z3=batch_z3,
+            shifts=batch_shifts,
+            N=N,
+            sampling_rate=1.0,
+        )
+        t_enter, t_exit, valid = self._compute_ray_box_intersections(origins, directions)
+
+        if getattr(self, "_cuda_graph_static_sampling", False):
+            return self._get_coords_box_fixed_ds_cuda_graph(
+                origins=origins,
+                directions=directions,
+                t_enter=t_enter,
+                t_exit=t_exit,
+                valid=valid,
+                num_samples_per_ray=num_samples_per_ray,
+            )
 
-        return transformed_rays
+        if not bool(valid.any()):
+            self._ray_meta = {
+                "valid": valid,
+                "local_ray_ids": torch.zeros(0, dtype=torch.long, device=origins.device),
+                "step_sizes": origins.new_zeros(0),
+                "num_valid": 0,
+            }
+            return origins.new_zeros((0, 3))
+
+        origins_v = origins[valid]
+        directions_v = directions[valid]
+        t_enter_v = t_enter[valid]
+        t_exit_v = t_exit[valid]
+
+        ds = self._resolve_ray_ds(num_samples_per_ray)
+        (
+            coords_v,
+            sample_mask_v,
+            n_samples_v,
+            lengths_v,
+        ) = self._sample_ray_segment_coords_fixed_ds(
+            origins_v, directions_v, t_enter_v, t_exit_v, ds
+        )
+
+        # nonzero() walks row-major, so local_ray_ids lines up with coords_v[sample_mask_v].
+        local_ray_ids = sample_mask_v.nonzero(as_tuple=False)[:, 0]
+        step_sizes_v = lengths_v / (n_samples_v.to(lengths_v.dtype) - 1.0)
+
+        self._ray_meta = {
+            "valid": valid,
+            "local_ray_ids": local_ray_ids,
+            "step_sizes": step_sizes_v,
+            "num_valid": int(origins_v.shape[0]),
+        }
+
+        all_coords = coords_v[sample_mask_v]
+        return all_coords.to(self.device, dtype=torch.float32)
+
+    def _get_coords_box_fixed_ds_cuda_graph(
+        self,
+        *,
+        origins: torch.Tensor,
+        directions: torch.Tensor,
+        t_enter: torch.Tensor,
+        t_exit: torch.Tensor,
+        valid: torch.Tensor,
+        num_samples_per_ray: int,
+    ) -> torch.Tensor:
+        """Build a fixed-capacity ray representation suitable for CUDA Graph replay.
+
+        Eager fixed-ds sampling compacts valid rays and sizes its padded dimension from a
+        GPU reduction. Both make allocation shapes depend on batch contents. During graph
+        capture, retain every ray and use the cube's maximum possible chord to choose one
+        conservative, host-known padded width. Invalid/padded coordinates are placed
+        outside the cube, so object-model masks make their densities and gradients zero.
+        """
+        ds = self._resolve_ray_ds(num_samples_per_ray)
+        max_samples = math.ceil((2.0 * math.sqrt(3.0)) / float(ds)) + 1
+        lengths = t_exit - t_enter
+        n_samples = torch.ceil(lengths.detach() / float(ds)).to(torch.long) + 1
+        n_samples = torch.clamp(n_samples, min=2, max=max_samples)
+
+        sample_idx = torch.arange(max_samples, device=origins.device).unsqueeze(0)
+        sample_mask = valid.unsqueeze(1) & (sample_idx < n_samples.unsqueeze(1))
+        denom = (n_samples - 1).to(origins.dtype).unsqueeze(1)
+        u = sample_idx.to(origins.dtype) / denom
+        t_vals = t_enter.unsqueeze(1) + lengths.unsqueeze(1) * u
+        coords = origins.unsqueeze(1) + t_vals.unsqueeze(2) * directions.unsqueeze(1)
+        coords = torch.where(sample_mask.unsqueeze(2), coords, coords.new_full((), 2.0))
+
+        step_sizes = lengths / (n_samples.to(lengths.dtype) - 1.0)
+        self._ray_meta = {
+            "valid": valid,
+            "sample_mask": sample_mask,
+            "step_sizes": step_sizes,
+            "num_valid_samples": sample_mask.sum(),
+        }
+        return coords.reshape(-1, 3).to(self.device, dtype=torch.float32)
+
+    def integrate_rays(
+        self, rays: torch.Tensor, num_samples_per_ray: int, target_values_len: int
+    ) -> torch.Tensor:
+        if getattr(self, "ray_sampling", "legacy") == "box_fixed_ds":
+            return self._integrate_rays_box_fixed_ds(rays, target_values_len)
+        return self._integrate_rays_legacy(rays, num_samples_per_ray, target_values_len)
 
     @staticmethod
     @torch.compile(mode="reduce-overhead")
-    def integrate_rays(
+    def _integrate_rays_legacy(
         rays: torch.Tensor, num_samples_per_ray: int, target_values_len: int
     ) -> torch.Tensor:
         ray_densities = rays.view(
@@ -597,6 +1229,54 @@ def integrate_rays(
 
         return predicted_values
 
+    def _integrate_rays_box_fixed_ds(
+        self, densities: torch.Tensor, target_values_len: int
+    ) -> torch.Tensor:
+        """Riemann-sum integrate ragged fixed-ds samples back to one value per ray.
+
+        ``densities`` are the INR outputs for the flattened, mask-selected samples that
+        ``get_coords`` produced (same order). Scatter-add them per ray and scale by the
+        per-ray physical step, writing into a zero (B,) output at the rays that actually
+        intersected the cube.
+        """
+        meta = self._ray_meta
+        if meta is None:
+            raise RuntimeError(
+                "integrate_rays called in 'box_fixed_ds' mode without ray metadata; "
+                "get_coords must run first."
+            )
+        if "sample_mask" in meta:
+            sample_mask = meta["sample_mask"]
+            ray_densities = densities.reshape(sample_mask.shape)
+            ray_sums = (ray_densities * sample_mask.to(densities.dtype)).sum(dim=1)
+            predicted = ray_sums * meta["step_sizes"].to(densities.dtype)
+            return torch.where(meta["valid"], predicted, torch.zeros_like(predicted))
+
+        predicted = torch.zeros(target_values_len, device=densities.device, dtype=densities.dtype)
+        valid = meta["valid"]
+        if not bool(valid.any()):
+            self._ray_meta = None
+            return predicted
+
+        local_ray_ids = meta["local_ray_ids"]
+        step_sizes = meta["step_sizes"].to(densities.dtype)
+        num_valid = int(meta["num_valid"])
+
+        ray_sums = torch.zeros(num_valid, device=densities.device, dtype=densities.dtype)
+        ray_sums.index_add_(0, local_ray_ids, densities)
+        predicted[valid] = ray_sums * step_sizes
+        # Consumed; clear so a stale batch can never be silently reused.
+        self._ray_meta = None
+        return predicted
+
+    def graph_constraint_densities(self, densities: torch.Tensor) -> torch.Tensor:
+        """Match eager sparsity normalization for graph-padded fixed-ds densities."""
+        meta = self._ray_meta
+        if meta is None or "sample_mask" not in meta:
+            return densities
+        num_valid_samples = meta["num_valid_samples"].clamp_min(1).to(densities.dtype)
+        return densities * (densities.numel() / num_valid_samples)
+
     # --- Torch Dataset Methods ---
     def __getitem__(
         self,
@@ -611,14 +1291,23 @@ def __getitem__(
         projection_idx = actual_idx // (self.tilt_stack.shape[1] * self.tilt_stack.shape[2])
         remaining = actual_idx % (self.tilt_stack.shape[1] * self.tilt_stack.shape[2])
 
-        pixel_i = remaining // self.tilt_stack.shape[1]
-        pixel_j = remaining % self.tilt_stack.shape[1]
+        pixel_i = remaining // self.tilt_stack.shape[2]
+        pixel_j = remaining % self.tilt_stack.shape[2]
+        if self.tilt_angles_per_row is not None:
+            phi = self.tilt_angles_per_row[projection_idx, pixel_i]
+        elif self.tilt_angles_per_col is not None:
+            phi = self.tilt_angles_per_col[projection_idx, pixel_j]
+        else:
+            phi = self.tilt_angles[projection_idx]
 
+        # Plain ints for the index fields: default_collate builds one int64 tensor per
+        # batch either way, but wrapping each index in torch.tensor() here allocates
+        # three scalar tensors per item on the dataloader hot path.
         return {
-            "projection_idx": torch.tensor(projection_idx),
-            "pixel_i": torch.tensor(pixel_i),
-            "pixel_j": torch.tensor(pixel_j),
-            "phi": self.tilt_angles[projection_idx],  # tensor
+            "projection_idx": projection_idx,
+            "pixel_i": pixel_i,
+            "pixel_j": pixel_j,
+            "phi": phi,  # tensor
             "target_value": self.tilt_stack[projection_idx, pixel_i, pixel_j],  # tensor
         }
 
@@ -628,13 +1317,10 @@ def __len__(
         """
         Returns the number of pixels in the tilt stack.
         """
-        N = max(self.tilt_stack.shape)
-        return self.tilt_stack.shape[0] * N * N
+        return self.tilt_stack.shape[0] * self.tilt_stack.shape[1] * self.tilt_stack.shape[2]
 
     def to(self, device: torch.device | str):
-        self._z1_params = nn.Parameter(self._z1_angles.to(device))
-        self._z3_params = nn.Parameter(self._z3_angles.to(device))
-        self._shifts_params = nn.Parameter(self._shifts.to(device))
+        self._materialize_pose_parameters(device)
 
         self._z1_ref = self._z1_ref.to(device)
         self._z3_ref = self._z3_ref.to(device)
@@ -691,6 +1377,13 @@ def __init__(
         else:
             data_quantile = torch.quantile(data, 0.95)
 
+        # Same guard as TomographyDatasetBase: a >95%-zero target has a zero
+        # 95th quantile and would normalize to inf/NaN.
+        if data_quantile <= 0:
+            data_quantile = data.abs().max()
+        if data_quantile <= 0:
+            raise ValueError("pretrain_target is all zeros; cannot normalize.")
+
         data = data / data_quantile
         data = torch.permute(data, (0, 3, 2, 1))
         # data = torch.flip(data, dims=(2,))
diff --git a/src/quantem/tomography/logger_tomography.py b/src/quantem/tomography/logger_tomography.py
index e07072f72..931a12aef 100644
--- a/src/quantem/tomography/logger_tomography.py
+++ b/src/quantem/tomography/logger_tomography.py
@@ -1,3 +1,5 @@
+from typing import Any, Literal, Mapping
+
 import matplotlib.pyplot as plt
 import numpy as np
 import torch
@@ -10,6 +12,11 @@
 class LoggerTomography(LoggerBase):
     """
     Logger for ML-based tomography reconstructions.
+
+    ``mode="tensorboard"`` preserves the original SummaryWriter behavior. ``mode="wandb"``
+    mirrors the same tags and steps to WandB, defaults to ``WANDB_MODE=offline`` unless the
+    environment already sets it, and writes offline run files under ``<log_dir>/wandb/``.
+    Upload offline runs with ``wandb sync <log_dir>/wandb/<offline-run-dir>``.
     """
 
     def __init__(
@@ -18,13 +25,30 @@ def __init__(
         run_prefix: str,
         run_suffix: str = "",
         log_images_every: int = 10,
+        mode: Literal["tensorboard", "wandb"] | str = "tensorboard",
+        wandb_config: Mapping[str, Any] | None = None,
     ):
-        super().__init__(log_dir, run_prefix, run_suffix, log_images_every)
+        super().__init__(
+            log_dir,
+            run_prefix,
+            run_suffix,
+            log_images_every,
+            mode=mode,
+            wandb_config=wandb_config,
+        )
 
-    def log_epoch(self, epoch: int, loss: float, tilt_series_loss: float, soft_loss: float):
-        self.log_scalar("loss/total", loss, epoch)
-        self.log_scalar("loss/tilt_series", tilt_series_loss, epoch)
-        self.log_scalar("loss/soft", soft_loss, epoch)
+    def log_epoch(
+        self,
+        epoch: int,
+        loss: float,
+        tilt_series_loss: float,
+        soft_loss: float,
+        grad_step: int | None = None,
+    ):
+        extra_steps = {"grad_step": grad_step} if grad_step is not None else None
+        self.log_scalar("loss/total", loss, epoch, extra_steps=extra_steps)
+        self.log_scalar("loss/tilt_series", tilt_series_loss, epoch, extra_steps=extra_steps)
+        self.log_scalar("loss/soft", soft_loss, epoch, extra_steps=extra_steps)
 
     def log_iter(
         self,
@@ -35,15 +59,31 @@ def log_iter(
         learning_rates: dict[str, float],
         num_samples_per_ray: int,
         val_loss: float | None = None,
+        val_fg_loss: float | None = None,
+        val_bg_loss: float | None = None,
+        grad_step: int | None = None,
     ):
-        self.log_scalar("loss/consistency", consistency_loss, iter)
-        self.log_scalar("loss/total", total_loss, iter)
-        self.log_scalar("loss/soft", object_model._soft_constraint_losses[-1], iter)
-        self.log_scalar("num_samples_per_ray", num_samples_per_ray, iter)
+        extra_steps = {"grad_step": grad_step} if grad_step is not None else None
+        self.log_scalar("loss/consistency", consistency_loss, iter, extra_steps=extra_steps)
+        self.log_scalar("loss/total", total_loss, iter, extra_steps=extra_steps)
+        self.log_scalar(
+            "loss/soft",
+            object_model._soft_constraint_losses[-1],
+            iter,
+            extra_steps=extra_steps,
+        )
+        self.log_scalar("num_samples_per_ray", num_samples_per_ray, iter, extra_steps=extra_steps)
         for param_name, lr_value in learning_rates.items():
-            self.log_scalar(f"learning_rate/{param_name}", float(lr_value), iter)
+            self.log_scalar(
+                f"learning_rate/{param_name}", float(lr_value), iter, extra_steps=extra_steps
+            )
         if val_loss is not None:
-            self.log_scalar("loss/val", val_loss, iter)
+            self.log_scalar("loss/validation", val_loss, iter, extra_steps=extra_steps)
+            self.log_scalar("loss/val", val_loss, iter, extra_steps=extra_steps)
+        if val_fg_loss is not None:
+            self.log_scalar("val/fg", val_fg_loss, iter, extra_steps=extra_steps)
+        if val_bg_loss is not None:
+            self.log_scalar("val/bg", val_bg_loss, iter, extra_steps=extra_steps)
 
     def log_iter_images(
         self,
diff --git a/src/quantem/tomography/object_models.py b/src/quantem/tomography/object_models.py
index 817e7ed5d..ef0dbecb7 100644
--- a/src/quantem/tomography/object_models.py
+++ b/src/quantem/tomography/object_models.py
@@ -1,4 +1,7 @@
+import os
+import weakref
 from abc import abstractmethod
+from contextlib import contextmanager
 from copy import deepcopy
 from dataclasses import dataclass
 from typing import Any, Callable, Generator, Optional, cast
@@ -9,15 +12,18 @@
 import torch.nn as nn
 from tqdm.auto import tqdm
 
+from quantem.core import config
 from quantem.core.io.serialize import AutoSerialize
 from quantem.core.ml.constraints import BaseConstraints, Constraints
 from quantem.core.ml.ddp import DDPMixin
 from quantem.core.ml.loss_functions import get_loss_module
 from quantem.core.ml.models.model_base import PlanarDecompositionModel
 from quantem.core.ml.optimizer_mixin import OptimizerMixin
+from quantem.core.ml.s3im import S3IMLoss
 from quantem.core.utils.rng import RNGMixin
 from quantem.tomography.dataset_models import TomographyINRPretrainDataset
 from quantem.tomography.tomography_context import ReconstructionContext
+from quantem.tomography.utils import tv_loss_vol_sq
 
 
 class ObjConstraintParams:
@@ -99,9 +105,15 @@ class ObjINRConstraints(Constraints):
         shrinkage: float = 0.0
         tv_vol: float = 0.0
         sparsity: float = 0.0
+        # S3IM (stochastic structural similarity) multiplex loss. s3im_weight is
+        # the penalty weight; the rest are config for the SSIM patch (paper defaults).
+        s3im_weight: float = 0.0
+        s3im_repeat_time: int = 10
+        s3im_kernel: int = 4
+        s3im_value_range: float = 1.0
         _name: str = "obj_inr"
 
-        soft_constraint_keys = ["tv_vol", "sparsity"]
+        soft_constraint_keys = ["tv_vol", "sparsity", "s3im_weight"]
         hard_constraint_keys = ["positivity", "shrinkage"]
 
     @dataclass
@@ -128,9 +140,20 @@ class ObjTensorDecompConstraints(Constraints):
         tv_vol: float = 0.0
         tv_plane: float = 0.0
         sparsity: float = 0.0
+        # Wedge-aware anisotropic volume TV: per-axis (z,y,x) weight multipliers applied to
+        # the volume-TV gradient norm. None => isotropic (original behaviour). Penalizing the
+        # missing-wedge (y/z) directions more than the resolved (x) direction suppresses
+        # streak/elongation artifacts without blurring the well-resolved in-plane structure.
+        tv_vol_aniso: tuple | None = None
+        # S3IM (stochastic structural similarity) multiplex loss. s3im_weight is
+        # the penalty weight; the rest are config for the SSIM patch (paper defaults).
+        s3im_weight: float = 0.0
+        s3im_repeat_time: int = 10
+        s3im_kernel: int = 4
+        s3im_value_range: float = 1.0
         _name: str = "obj_tensor_decomp"
 
-        soft_constraint_keys = ["tv_vol", "tv_plane", "sparsity"]
+        soft_constraint_keys = ["tv_vol", "tv_plane", "sparsity", "s3im_weight"]
         hard_constraint_keys = ["positivity", "shrinkage"]
 
     @classmethod
@@ -204,6 +227,48 @@ def _unwrap(model: nn.Module | nn.parallel.DistributedDataParallel) -> PlanarDec
     return cast(PlanarDecompositionModel, model)
 
 
+def _plane_tv_loss_eager(
+    grids: torch.nn.ParameterList | list[torch.Tensor] | tuple[torch.Tensor, ...],
+    tilted: bool,
+    rotations: int,
+) -> torch.Tensor:
+    """Reference plane-TV reduction retained for CPU and capability fallback."""
+    per_level = []
+    for plane in grids:
+        dh = (plane[:, :, 1:, :] - plane[:, :, :-1, :]).pow(2).mean(dim=(1, 2, 3))
+        dw = (plane[:, :, :, 1:] - plane[:, :, :, :-1]).pow(2).mean(dim=(1, 2, 3))
+        per_plane = dh + dw
+        if tilted:
+            per_rotation = per_plane.view(rotations, 3).sum(dim=1)
+            level_tv = per_rotation.mean()
+        else:
+            level_tv = per_plane.sum()
+        per_level.append(level_tv)
+    return torch.stack(per_level).sum()
+
+
+def _plane_tv_loss(
+    grids: torch.nn.ParameterList | list[torch.Tensor] | tuple[torch.Tensor, ...],
+    tilted: bool,
+    rotations: int,
+) -> torch.Tensor:
+    """Dispatch three fp32 CUDA levels to quantem-cuda when available."""
+    use_fused = (
+        os.environ.get("QUANTEM_PLANE_TV_FUSED", "1") != "0"
+        and len(grids) == 3
+        and all(grid.is_cuda and grid.dtype == torch.float32 for grid in grids)
+        and config.get("has_quantem_cuda")
+        and config.get("use_cuda_kernels", default=True)
+    )
+    if use_fused:
+        import quantem.cuda.core.ml as cuda_ml
+
+        fused = getattr(cuda_ml, "plane_tv_loss", None)
+        if fused is not None:
+            return fused(grids[0], grids[1], grids[2])
+    return _plane_tv_loss_eager(grids, tilted, rotations)
+
+
 class ObjectBase(AutoSerialize, nn.Module, RNGMixin, OptimizerMixin):
     DEFAULT_LRS = {
         "object": 8e-6,
@@ -421,12 +486,10 @@ def apply_hard_constraints(
 
     def apply_soft_constraints(self, ctx: ReconstructionContext) -> torch.Tensor:
         assert ctx.obj is not None, "ObjectPixelated requires ctx.obj to be set"
-        soft_loss = torch.tensor(
-            0.0, device=ctx.obj.device, dtype=ctx.obj.dtype, requires_grad=True
-        )
+        soft_loss = torch.tensor(0.0, device=ctx.obj.device, dtype=ctx.obj.dtype)
         if self.constraints.tv_vol > 0:
             tv_loss = self.get_tv_loss(ctx)
-            soft_loss += tv_loss
+            soft_loss = soft_loss + tv_loss
         return soft_loss
 
     # --- Forward method ---
@@ -439,11 +502,9 @@ def get_tv_loss(self, ctx: ReconstructionContext) -> torch.Tensor:
         # TV over the three trailing spatial dims, leaving any leading channel/batch axes
         # intact. Works for a 3-D volume, obj_view's [1, D, H, W], and a multimodal
         # [C, D, H, W] (channels = elemental compositions), matching the INR / tensor-decomp
-        # convention where the object carries a leading channel dimension.
-        tv_d = torch.pow(ctx.obj[..., 1:, :, :] - ctx.obj[..., :-1, :, :], 2).sum()
-        tv_h = torch.pow(ctx.obj[..., :, 1:, :] - ctx.obj[..., :, :-1, :], 2).sum()
-        tv_w = torch.pow(ctx.obj[..., :, :, 1:] - ctx.obj[..., :, :, :-1], 2).sum()
-        tv_loss = tv_d + tv_h + tv_w
+        # convention where the object carries a leading channel dimension. tv_loss_vol_sq
+        # dispatches to the fused quantem-cuda kernel when available.
+        tv_loss = tv_loss_vol_sq(ctx.obj)
 
         return tv_loss * self.constraints.tv_vol / ctx.obj.numel()
 
@@ -457,6 +518,14 @@ def to(self, device: str | torch.device):
         return self
 
 
+# torch.compile artifacts keyed by the live model object. Kept outside the
+# instances so AutoSerialize never sees them and reset()/rebuild_model()
+# (which swap the model object) naturally invalidate the cache.
+_compiled_forward_cache: "weakref.WeakKeyDictionary[nn.Module, Callable]" = (
+    weakref.WeakKeyDictionary()
+)
+
+
 class ObjectINR(ObjectConstraints, DDPMixin):
     DEFAULT_CONSTRAINTS = ObjConstraintParams.ObjINRConstraints()
 
@@ -466,6 +535,7 @@ def __init__(
         device: str = "cpu",
         rng: np.random.Generator | int | None = None,
         model: nn.Module | None = None,
+        compile_model: bool = False,
         _token: object | None = None,
     ):
         super().__init__(
@@ -476,6 +546,7 @@ def __init__(
         )
         self._pretrain_losses = []
         self._pretrain_lrs = []
+        self._compile_model = bool(compile_model)
         self.constraints: ObjConstraintParams.ObjINRConstraints = self.DEFAULT_CONSTRAINTS.copy()
         # Register the network submodule (important: real nn.Module attribute)
         if model is not None:
@@ -489,18 +560,51 @@ def from_model(
         shape: tuple[int, int, int],
         device: str = "cpu",
         rng: np.random.Generator | int | None = None,
+        compile_model: bool = False,
     ):
         obj_model = cls(
             shape=shape,
             device=device,
             rng=rng,
             model=model,  # ✅ build/register in __init__
+            compile_model=compile_model,
         )
 
         obj_model.setup_distributed(device=device)
         obj_model.to(device)
         return obj_model
 
+    def _prepare_model_call(self) -> None:
+        """Hook for subclasses that select optional model capabilities."""
+
+    @contextmanager
+    def reconstruction_forward_context(self):
+        """Scope optional auxiliaries to a reconstruction's primary model call."""
+        yield
+
+    def _unpack_model_output(self, output: Any) -> torch.Tensor:
+        """Return the primary tensor from models with auxiliary outputs."""
+        if isinstance(output, tuple):
+            return output[0]
+        return output
+
+    def _model_call(self, coords: torch.Tensor) -> Any:
+        """Invoke the model, through torch.compile when compile_model was set.
+
+        Compiles the bound __call__ (a plain function, so nothing extra is
+        registered on the module tree or picked up by AutoSerialize) and
+        caches per model object.
+        """
+        self._prepare_model_call()
+        model = self.model
+        if not getattr(self, "_compile_model", False):
+            return model(coords)
+        fn = _compiled_forward_cache.get(model)
+        if fn is None:
+            fn = torch.compile(model.__call__)
+            _compiled_forward_cache[model] = fn
+        return fn(coords)
+
     # --- Properties ---
 
     @property
@@ -542,7 +646,8 @@ def apply_soft_constraints(
         self,
         ctx: ReconstructionContext,
     ) -> torch.Tensor:
-        soft_loss = torch.tensor(0.0, device=ctx.coords.device)
+        device = ctx.coords.device if ctx.coords is not None else self._device
+        soft_loss = torch.zeros((), device=device)
         if self.constraints.tv_vol > 0:
             assert ctx.coords is not None, (
                 "coords must be provided for INR object model to compute the TV loss"
@@ -559,6 +664,9 @@ def apply_soft_constraints(
             sparsity_loss = self.constraints.sparsity * torch.norm(ctx.pred, p=1)
             soft_loss += sparsity_loss
 
+        if getattr(self.constraints, "s3im_weight", 0.0) > 0:
+            soft_loss += self.get_s3im_loss(ctx)
+
         return soft_loss
 
     def apply_hard_constraints(self, pred: torch.Tensor) -> torch.Tensor:
@@ -573,6 +681,10 @@ def apply_hard_constraints(self, pred: torch.Tensor) -> torch.Tensor:
 
         return pred
 
+    def sample_tv_tap_coords(self, coords: torch.Tensor) -> Optional[torch.Tensor]:
+        """Hook for the training loop: returns None (INR TV uses autograd, no tap merging)."""
+        return None
+
     # --- Define get_tv_loss ---
 
     def get_tv_loss(self, ctx: ReconstructionContext) -> torch.Tensor:
@@ -604,6 +716,33 @@ def get_tv_loss(self, ctx: ReconstructionContext) -> torch.Tensor:
         grad_norm = torch.norm(grad_outputs, dim=1)  # Shape: [num_samples]
         return self.constraints.tv_vol * grad_norm.mean()
 
+    # --- Define get_s3im_loss ---
+
+    def get_s3im_loss(self, ctx: ReconstructionContext) -> torch.Tensor:
+        """
+        Compute the weighted S3IM (stochastic structural similarity) multiplex loss
+        between the predicted and measured projection pixels of the batch.
+
+        S3IM consumes the already-computed ``ctx.pred`` (kept attached to the graph)
+        and ``ctx.target`` -- it does NOT re-evaluate the model, so it adds no extra
+        forward pass and is covered by the single backward of the reconstruction step.
+        """
+        assert ctx.pred is not None, "pred must be provided to compute the S3IM loss"
+        assert ctx.target is not None, "target must be provided to compute the S3IM loss"
+
+        # Lazily build the (paramless) SSIM module and cache it without registering it
+        # as an nn.Module submodule (avoids perturbing state_dict / AutoSerialize / DDP).
+        s3im = getattr(self, "_s3im", None)
+        if s3im is None:
+            s3im = S3IMLoss(
+                repeat_time=self.constraints.s3im_repeat_time,
+                kernel_size=self.constraints.s3im_kernel,
+                value_range=self.constraints.s3im_value_range,
+            ).to(ctx.pred.device)
+            object.__setattr__(self, "_s3im", s3im)
+
+        return self.constraints.s3im_weight * s3im(ctx.pred, ctx.target)
+
     # --- Optimization Parameters ---
     @property
     def params(self) -> Generator[torch.nn.Parameter, None, None]:
@@ -657,12 +796,20 @@ def forward(self, coords: Optional[torch.Tensor] = None) -> torch.Tensor:
         """forward pass for the INR model"""
         assert coords is not None, "ObjectINR.forward requires coords"
 
-        all_densities = self.model(coords)
+        # TV's autograd.grad recompute stays on the eager model (double
+        # backward through compiled graphs is not reliable); only this main
+        # forward goes through the compiled path.
+        all_densities = self._unpack_model_output(self._model_call(coords))
 
         if all_densities.dim() > 1:
             all_densities = all_densities.squeeze(-1)
         valid_mask = (
-            (coords[:, 0] >= -1) & (coords[:, 0] <= 1) & (coords[:, 1] >= -1) & (coords[:, 1] <= 1)
+            (coords[:, 0] >= -1)
+            & (coords[:, 0] <= 1)
+            & (coords[:, 1] >= -1)
+            & (coords[:, 1] <= 1)
+            & (coords[:, 2] >= -1)
+            & (coords[:, 2] <= 1)
         ).float()
 
         if all_densities.dim() > 1:
@@ -674,6 +821,39 @@ def forward(self, coords: Optional[torch.Tensor] = None) -> torch.Tensor:
 
         return all_densities
 
+    def forward_with_tv_taps(
+        self, coords: torch.Tensor, tap_coords: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Single model call covering the main batch and the volume-TV tap points.
+
+        The out-of-bounds mask and hard constraints apply to the main chunk
+        only; tap densities are returned raw (border-clamped), matching the
+        fallback path in ``get_volume_tv_loss``.
+        """
+        self._prepare_model_call()
+        merged = self._unpack_model_output(self.model(torch.cat([coords, tap_coords], dim=0)))
+        main, taps = merged[: coords.shape[0]], merged[coords.shape[0] :]
+
+        if main.dim() > 1:
+            main = main.squeeze(-1)
+        valid_mask = (
+            (coords[:, 0] >= -1)
+            & (coords[:, 0] <= 1)
+            & (coords[:, 1] >= -1)
+            & (coords[:, 1] <= 1)
+            & (coords[:, 2] >= -1)
+            & (coords[:, 2] <= 1)
+        ).float()
+        if main.dim() > 1:
+            valid_mask = valid_mask.unsqueeze(-1)
+        main = main * valid_mask
+        main = self.apply_hard_constraints(main)
+
+        if taps.dim() == 1:
+            taps = taps.unsqueeze(-1)
+        return main, taps
+
     # Pretrain Loop
 
     def pretrain(
@@ -774,7 +954,7 @@ def create_volume(self, return_vol: bool = False):
             coords_1d = torch.linspace(-1, 1, N)
             x, y, z = torch.meshgrid(coords_1d, coords_1d, coords_1d, indexing="ij")
             inputs = torch.stack([x, y, z], dim=-1).reshape(-1, 3)
-            model = self.model.module if isinstance(self.model, nn.DataParallel) else self.model
+            model = _unwrap(self.model)
 
             inference_batch_size = 5 * N * N
             total_samples = N**3
@@ -874,12 +1054,14 @@ def __init__(
         device: str = "cpu",
         rng: np.random.Generator | int | None = None,
         model: nn.Module | None = None,
+        compile_model: bool = False,
         _token: object | None = None,
     ):
         super().__init__(
             shape=shape,
             device=device,
             rng=rng,
+            compile_model=compile_model,
             _token=self._token,
         )
         self._pretrain_losses = []
@@ -899,25 +1081,91 @@ def from_model(
         shape: tuple[int, int, int],
         device: str = "cpu",
         rng: np.random.Generator | int | None = None,
+        compile_model: bool = False,
     ):
         obj_model = cls(
             shape=shape,
             device=device,
             rng=rng,
             model=model,  # ✅ build/register in __init__
+            compile_model=compile_model,
         )
 
         obj_model.setup_distributed(device=device)
         obj_model.to(device)
         return obj_model
 
+    def _prepare_model_call(self) -> None:
+        model = _unwrap(self.model)
+        use_combined = (
+            getattr(self, "_reconstruction_plane_tv_requested", False)
+            and os.environ.get("QUANTEM_KPLANES_MS_TV_FUSED", "1") != "0"
+            and self.constraints.tv_plane > 0
+            and getattr(model, "quantem_supports_fused_plane_tv", False)
+            and config.get("has_quantem_cuda")
+            and config.get("use_cuda_kernels", default=True)
+        )
+        if use_combined:
+            try:
+                import quantem.cuda.core.ml as cuda_ml
+            except (ImportError, OSError, RuntimeError):
+                use_combined = False
+            else:
+                use_combined = getattr(cuda_ml, "kplanes_tilted_fuse_ms_tv", None) is not None
+        model._plane_tv_fusion_requested = use_combined
+        self._fused_plane_tv_loss = None
+
+    @contextmanager
+    def reconstruction_forward_context(self):
+        """Request the fused plane-TV auxiliary for one reconstruction forward."""
+        self._reconstruction_plane_tv_requested = True
+        completed = False
+        try:
+            yield
+            completed = True
+        finally:
+            self._reconstruction_plane_tv_requested = False
+            _unwrap(self.model)._plane_tv_fusion_requested = False
+            if not completed:
+                self._fused_plane_tv_loss = None
+
+    def _unpack_model_output(self, output: Any) -> torch.Tensor:
+        if isinstance(output, tuple):
+            density = output[0]
+            if len(output) == 2 and isinstance(output[1], torch.Tensor) and output[1].ndim == 0:
+                self._fused_plane_tv_loss = output[1]
+            return density
+        return output
+
+    def sample_tv_tap_coords(self, coords: torch.Tensor) -> Optional[torch.Tensor]:
+        """
+        Sample the finite-difference tap coordinates for the volume TV loss.
+
+        Returns a (4*n, 3) tensor [base; base+h*ex; base+h*ey; base+h*ez] for n
+        sampled base points, or None when tv_vol == 0.  The training loop
+        concatenates this to all_coords so the TV taps are evaluated in the
+        same model call as the main forward pass.
+        """
+        if self.constraints.tv_vol == 0:
+            return None
+        model = _unwrap(self.model)
+        h = 2.0 / min(model.resolution)
+        num_tv_samples = min(10_000, coords.shape[0])
+        tv_indices = torch.randperm(coords.shape[0], device=coords.device)[:num_tv_samples]
+        tv_coords = coords[tv_indices]  # (n, 3)
+        offsets = h * torch.eye(3, device=tv_coords.device)
+        return torch.cat(
+            [tv_coords, tv_coords + offsets[0], tv_coords + offsets[1], tv_coords + offsets[2]],
+            dim=0,
+        )
+
     # --- Constraints ---
 
     def apply_soft_constraints(self, ctx: ReconstructionContext) -> torch.Tensor:
-        soft_loss = torch.tensor(
-            0.0, device=ctx.pred.device if ctx.pred is not None else self.device
+        soft_loss = torch.zeros(
+            (), device=ctx.pred.device if ctx.pred is not None else self.device
         )
-        if self.constraints.tv_vol > 0:
+        if self.constraints.tv_vol > 0 or self.constraints.tv_plane > 0:
             assert ctx.coords is not None, "Coordinates must be provided for TV loss"
             assert ctx.pred is not None, "Prediction must be provided for TV loss"
             soft_loss += self.get_tv_loss(ctx)
@@ -929,6 +1177,9 @@ def apply_soft_constraints(self, ctx: ReconstructionContext) -> torch.Tensor:
             sparsity_loss = self.constraints.sparsity * ctx.all_densities.abs().mean()
             soft_loss += sparsity_loss
 
+        if self.constraints.s3im_weight > 0:
+            soft_loss += self.get_s3im_loss(ctx)
+
         return soft_loss
 
     # TV Losses
@@ -942,68 +1193,92 @@ def get_tv_loss(self, ctx: ReconstructionContext) -> torch.Tensor:
         """
         assert ctx.coords is not None, "Coordinates must be provided for TV loss"
         assert ctx.pred is not None, "Prediction must be provided for TV loss"
-        tv_loss = torch.tensor(0.0, device=ctx.pred.device)
-        tv_loss += self._get_plane_tv_loss()
-        tv_loss += self.get_volume_tv_loss(ctx.coords)
+        tv_loss = torch.zeros((), device=ctx.pred.device)
+        if self.constraints.tv_plane > 0:
+            fused_plane_tv = getattr(self, "_fused_plane_tv_loss", None)
+            if fused_plane_tv is None:
+                tv_loss = tv_loss + self._get_plane_tv_loss()
+            else:
+                tv_loss = tv_loss + self.constraints.tv_plane * fused_plane_tv
+                self._fused_plane_tv_loss = None
+                _unwrap(self.model)._plane_tv_fusion_requested = False
+        if self.constraints.tv_vol > 0:
+            tv_loss = tv_loss + self.get_volume_tv_loss(
+                ctx.coords, precomputed_tap_densities=ctx.tv_tap_densities
+            )
         return tv_loss
 
     def _get_plane_tv_loss(self) -> torch.Tensor:
         """
         Gets the total-variation across the planes.
         """
-        is_tilted = self.model.tilted
-        per_level = []
-
         model = _unwrap(self.model)
-        for p in model.grids:
-            # p: (3*T, C, H, W) for TILTED, (3, C, H, W) for KPlanes
-            dh = (p[:, :, 1:, :] - p[:, :, :-1, :]).pow(2).mean(dim=(1, 2, 3))
-            dw = (p[:, :, :, 1:] - p[:, :, :, :-1]).pow(2).mean(dim=(1, 2, 3))
-            per_plane = dh + dw  # (3*T,) or (3,)
-
-            if is_tilted:
-                T = self.model.T
-                per_rotation = per_plane.view(T, 3).sum(dim=1)  # sum 3 planes per rotation
-                level_tv = per_rotation.mean()  # avg across rotations
-            else:
-                level_tv = per_plane.sum()
-
-            per_level.append(level_tv)
+        rotations = model.T if model.tilted else 1
+        return self.constraints.tv_plane * _plane_tv_loss(model.grids, model.tilted, rotations)
 
-        return self.constraints.tv_plane * torch.stack(per_level).sum()
-
-    def get_volume_tv_loss(self, coords: torch.Tensor) -> torch.Tensor:
+    def get_volume_tv_loss(
+        self,
+        coords: torch.Tensor,
+        precomputed_tap_densities: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
         """
         Isotropic volume TV via finite differences. Same form as the autograd
         version (L1 of gradient L2-norm) but avoids double-backward, so it
         works for KPlanesTILTED, CPTilted, and anything else.
-        """
-        num_tv_samples = min(10_000, coords.shape[0])
-        tv_indices = torch.randperm(coords.shape[0], device=coords.device)[:num_tv_samples]
-        tv_coords = coords[tv_indices]  # (N, 3)
 
+        When *precomputed_tap_densities* is provided (a (4N, C) tensor from the
+        merged single-pass forward in the training loop), the model call is
+        skipped entirely and the supplied values are used directly.  When absent
+        the existing batched 4N-point fallback path runs unchanged.
+        """
         model = _unwrap(self.model)
         h = 2.0 / min(model.resolution)
 
-        pred = model(tv_coords)
-        if isinstance(pred, tuple):
-            pred = pred[0]
-        if pred.dim() == 1:
-            pred = pred.unsqueeze(-1)  # (N, 1)
-
-        grads = []
-        for axis in range(3):
-            offset = torch.zeros(3, device=tv_coords.device)
-            offset[axis] = h
-            shifted_pred = self.model(tv_coords + offset)
-            if isinstance(shifted_pred, tuple):
-                shifted_pred = shifted_pred[0]
-            if shifted_pred.dim() == 1:
-                shifted_pred = shifted_pred.unsqueeze(-1)
-            grads.append((shifted_pred - pred) / h)  # (N, 1)
-
-        grad_stack = torch.stack(grads, dim=-1)  # (N, C, 3)
-        grad_norm = torch.norm(grad_stack, dim=-1)  # (N, C)
+        if precomputed_tap_densities is not None:
+            # (4N, C) tap densities from the merged single-pass forward in the
+            # training loop; layout [base; +ex; +ey; +ez] per sample_tv_tap_coords.
+            all_pred = precomputed_tap_densities
+        else:
+            num_tv_samples = min(10_000, coords.shape[0])
+            tv_indices = torch.randperm(coords.shape[0], device=coords.device)[:num_tv_samples]
+            tv_coords = coords[tv_indices]  # (N, 3)
+
+            # Evaluate the base points and the three axis-shifted copies in a single
+            # batched forward (4N points) instead of 4 sequential model calls.
+            offsets = h * torch.eye(3, device=tv_coords.device, dtype=tv_coords.dtype)  # (3, 3)
+            all_coords = torch.cat(
+                [tv_coords, (tv_coords.unsqueeze(0) + offsets.unsqueeze(1)).reshape(-1, 3)]
+            )  # (4N, 3)
+
+            all_pred = model(all_coords)
+            if isinstance(all_pred, tuple):
+                all_pred = all_pred[0]
+
+        if all_pred.dim() == 1:
+            all_pred = all_pred.unsqueeze(-1)  # (4N, 1)
+
+        n = all_pred.shape[0] // 4
+        pred = all_pred[:n]  # (N, C)
+        shifted_pred = all_pred[n:].view(3, n, -1)  # (3, N, C)
+
+        grad_stack = (
+            shifted_pred - pred.unsqueeze(0)
+        ) / h  # (3, N, C); axis0 = coord(0,1,2)=(x,y,z)
+
+        aniso = getattr(self.constraints, "tv_vol_aniso", None)
+        if aniso is not None:
+            # Weight per coordinate direction (coord0=x, coord1=y, coord2=z). The missing
+            # wedge lives in y/z, so set those weights high and x low to suppress streak
+            # smear without over-smoothing the resolved in-plane (x) structure.
+            w = torch.stack(
+                [
+                    torch.full((), value, device=grad_stack.device, dtype=grad_stack.dtype)
+                    for value in aniso
+                ]
+            ).view(3, 1, 1)
+            grad_stack = grad_stack * w
+
+        grad_norm = torch.norm(grad_stack, dim=0)  # (N, C)
 
         return self.constraints.tv_vol * grad_norm.mean()
 
@@ -1012,7 +1287,14 @@ def apply_hard_constraints(self, pred: torch.Tensor) -> torch.Tensor:
         Apply hard constraints to the predicted values of the INR model.
         """
 
-        if self.constraints.positivity:
+        activation = getattr(_unwrap(self.model), "density_activation", None)
+        activation_is_nonnegative = bool(
+            getattr(activation, "quantem_guarantees_nonnegative", False)
+        )
+        skip_redundant_positivity = (
+            os.environ.get("QUANTEM_DENSITY_TAIL_FUSED", "1") != "0" and activation_is_nonnegative
+        )
+        if self.constraints.positivity and not skip_redundant_positivity:
             pred = torch.clamp(pred, min=0.0, max=None)
         if self.constraints.shrinkage:
             pred = torch.max(pred - self.constraints.shrinkage, torch.zeros_like(pred))
diff --git a/src/quantem/tomography/tomography.py b/src/quantem/tomography/tomography.py
index f86095ade..2cfa8e6ac 100644
--- a/src/quantem/tomography/tomography.py
+++ b/src/quantem/tomography/tomography.py
@@ -1,6 +1,7 @@
 import os
+import warnings
 from pathlib import Path
-from typing import Literal, Self, Sequence
+from typing import Callable, Literal, Self, Sequence
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -10,15 +11,18 @@
 
 from quantem.core.io.serialize import load as autoserialize_load
 from quantem.core.ml.loss_functions import get_loss_module
-from quantem.core.ml.models.kplanes import CPTilted
+from quantem.core.ml.models.kplanes import CPTilted, KPlanesTILTED
+from quantem.core.ml.profiling import nsys_capture_tick
 from quantem.core.utils.filter import gaussian_filter_2d_stack, gaussian_kernel_1d
 from quantem.core.utils.tomography_utils import torch_phase_cross_correlation
 from quantem.tomography.dataset_models import (
     DatasetConstraintParams,
     DatasetConstraintsType,
     DatasetModelType,
+    DeviceBatchSampler,
     TomographyINRDataset,
     TomographyPixDataset,
+    build_pixel_holdout_split,
 )
 from quantem.tomography.logger_tomography import LoggerTomography
 from quantem.tomography.object_models import (
@@ -33,6 +37,56 @@
 from quantem.tomography.tomography_context import ReconstructionContext
 from quantem.tomography.tomography_opt import TomographyOpt
 
+if torch.cuda.is_available():
+    from torch.cuda import nvtx
+else:
+
+    class _NvtxNoop:
+        @staticmethod
+        def range_push(msg: str) -> None:
+            pass
+
+        @staticmethod
+        def range_pop() -> None:
+            pass
+
+    nvtx = _NvtxNoop()
+
+
+def _should_take_grad_step_snapshot(grad_step: int, snapshot_every: int) -> bool:
+    """Return whether a gradient-update snapshot should fire at ``grad_step``."""
+    return snapshot_every > 0 and grad_step > 0 and grad_step % snapshot_every == 0
+
+
+def _take_grad_step_snapshot(
+    *,
+    obj_model: ObjectINR | ObjectTensorDecomp,
+    grad_step: int,
+    global_rank: int,
+    logger: LoggerTomography | None,
+    snapshot_dir: str | Path | None,
+    snapshot_callback: Callable[[int, np.ndarray | None], None] | None,
+) -> None:
+    volume = obj_model.obj_view
+    volume_or_none = volume if global_rank == 0 else None
+
+    if global_rank == 0:
+        if snapshot_dir is not None:
+            snapshot_path = Path(snapshot_dir) / f"step_{grad_step}.npy"
+            snapshot_path.parent.mkdir(parents=True, exist_ok=True)
+            np.save(snapshot_path, np.asarray(volume_or_none).astype(np.float32, copy=False))
+
+        if logger is not None:
+            logger.log_scalar(
+                "snapshots/last_grad_step",
+                float(grad_step),
+                grad_step,
+                step_domain="grad_step",
+            )
+
+    if snapshot_callback is not None:
+        snapshot_callback(grad_step, volume_or_none)
+
 
 class Tomography(TomographyOpt, TomographyBase):
     """
@@ -60,6 +114,24 @@ def from_models(
             _token=cls._token,
         )
 
+    def _sync_pose_gradients_ddp(self) -> None:
+        """Average otherwise-unsynced pose gradients that live outside DDP.
+
+        Averaging rank-local shard gradients gives the full-batch direction and keeps
+        pose updates in lockstep across ranks.
+        """
+        if self.world_size <= 1 or not self.dset.has_optimizer():
+            return
+        if not dist.is_available() or not dist.is_initialized():
+            return
+
+        groups = self.dset.get_optimization_parameters()
+        for key in sorted(groups.keys()):
+            for p in groups[key]:
+                if p.grad is None:
+                    p.grad = torch.zeros_like(p)
+                dist.all_reduce(p.grad, op=dist.ReduceOp.AVG)
+
     def reconstruct(
         self,
         num_iter: int = 10,
@@ -73,6 +145,9 @@ def reconstruct(
         num_samples_per_ray: int | list[tuple[int, int]] | None = None,
         profiling_mode: bool = False,
         val_fraction: float = 0.0,
+        holdout_fraction: float = 0.0,
+        holdout_seed: int = 0,
+        holdout_every: int = 1,
         loss_type: Literal[
             "l2",
             "l1",
@@ -84,19 +159,86 @@ def reconstruct(
         loss_func_kwargs: dict = {},
         reset_dset: DatasetModelType | None = None,
         show_metrics: bool = False,
+        eval_callback: Callable[[int], None] | None = None,
+        eval_every: int = 0,
+        snapshot_every: int = 0,
+        snapshot_dir: str | Path | None = None,
+        snapshot_callback: Callable[[int, np.ndarray | None], None] | None = None,
+        pose_warmup_epochs: int = 0,
+        *,
+        autocast_dtype: str | torch.dtype | None = None,
+        grad_scaler: bool | None = None,
+        cuda_graphs: bool = False,
+        grad_clip_max_norm: float | None = 1.0,
     ):
         """
         This function should be able to handle both AD and INR-based tomography reconstruction methods.
         I.e, auto-detection through the obj model type, while both share the same pose optimization.
         """
+        if snapshot_every < 0:
+            raise ValueError("snapshot_every must be >= 0.")
+        if pose_warmup_epochs < 0:
+            raise ValueError("pose_warmup_epochs must be >= 0.")
+        if holdout_every < 1:
+            raise ValueError("holdout_every must be >= 1.")
+        if grad_clip_max_norm is not None and grad_clip_max_norm < 0:
+            raise ValueError("grad_clip_max_norm must be >= 0 or None.")
+        if val_fraction > 0.0 and holdout_fraction > 0.0:
+            raise ValueError("Use either val_fraction or holdout_fraction, not both.")
+        torch_dtype_names = {
+            torch.bfloat16: "bf16",
+            torch.float16: "fp16",
+            torch.float32: None,
+        }
+        if isinstance(autocast_dtype, torch.dtype):
+            autocast_dtype = torch_dtype_names.get(autocast_dtype, autocast_dtype)
+
+        autocast_dtypes = {"bf16": torch.bfloat16, "fp16": torch.float16}
+        if autocast_dtype is not None and (
+            not isinstance(autocast_dtype, str) or autocast_dtype not in autocast_dtypes
+        ):
+            raise ValueError(
+                "autocast_dtype must be one of None, 'bf16', 'fp16', "
+                "torch.bfloat16, torch.float16, or torch.float32."
+            )
+        grad_scaler_enabled = autocast_dtype == "fp16" if grad_scaler is None else grad_scaler
+        if grad_scaler_enabled and autocast_dtype is None:
+            raise ValueError("grad_scaler=True requires autocast_dtype to be set.")
+        autocast_enabled = autocast_dtype is not None
+        torch_autocast_dtype = autocast_dtypes.get(autocast_dtype, torch.bfloat16)
+        grad_scaler = torch.amp.GradScaler(self.device.type, enabled=grad_scaler_enabled)
+        snapshots_enabled = snapshot_every > 0
+
+        if cuda_graphs and self.world_size > 1:
+            raise NotImplementedError("CUDA Graph reconstruction does not support DDP.")
+        if cuda_graphs and torch.device(self.device).type != "cuda":
+            raise ValueError("cuda_graphs=True requires a CUDA device.")
+
+        # Reconnecting an existing optimizer during ``to`` marks all of its parameters
+        # trainable. Preserve an explicitly frozen tilted rotation bank; optimizer_params
+        # supplied below may intentionally enable it again.
+        tilted_so3_was_frozen = isinstance(self.obj_model.model, KPlanesTILTED) and not any(
+            parameter.requires_grad for parameter in self.obj_model.model.so3.parameters()
+        )
 
         # Check device consistency
         self.obj_model.to(self.device)
+        if tilted_so3_was_frozen:
+            self.obj_model.model.so3.requires_grad_(False)
+
+        previous_batch_size = getattr(self, "batch_size", None)
+        previous_num_workers = getattr(self, "num_workers", None)
+        previous_val_fraction = getattr(self, "val_fraction", None)
+        previous_holdout_fraction = getattr(self, "holdout_fraction", None)
+        previous_holdout_seed = getattr(self, "holdout_seed", None)
 
-        # Saving batch size, num workers, and val fraction for reloading
+        # Saving batch size, num workers, and validation split settings for reloading
         self.batch_size = batch_size
         self.num_workers = num_workers
         self.val_fraction = val_fraction
+        self.holdout_fraction = holdout_fraction
+        self.holdout_seed = holdout_seed
+        self.holdout_every = holdout_every
 
         if profiling_mode:
             if self.global_rank == 0:
@@ -106,10 +248,17 @@ def reconstruct(
             raise NotImplementedError("Reset is not implemented yet.")
 
         new_scheduler = reset
+        previous_cuda_graphs_optimizer = getattr(self.obj_model, "_cuda_graphs_optimizer", False)
+        if cuda_graphs:
+            # Optimizers constructed below use capturable fused Adam and device LR
+            # tensors. Existing optimizers remain usable for Phase A, but Phase B is
+            # enabled only if they were already built with these properties.
+            self.obj_model._cuda_graphs_optimizer = True
         if optimizer_params is not None:
             self.optimizer_params = optimizer_params
             self.set_optimizers()
             new_scheduler = True
+        self.obj_model._cuda_graphs_optimizer = previous_cuda_graphs_optimizer
 
         if scheduler_params is not None:
             self.scheduler_params = scheduler_params
@@ -129,8 +278,18 @@ def reconstruct(
                 dset_constraints = DatasetConstraintParams.parse_dict(dset_constraints)
 
             self.dset.constraints = dset_constraints
+        dataloader_needs_rebuild = (
+            not hasattr(self, "dataloader")
+            or reset_dset is not None
+            or previous_batch_size != batch_size
+            or previous_num_workers != num_workers
+            or previous_val_fraction != val_fraction
+            or previous_holdout_fraction != holdout_fraction
+            or previous_holdout_seed != holdout_seed
+        )
+
         # Setting up DDP
-        if not hasattr(self, "dataloader") or reset_dset is not None:
+        if dataloader_needs_rebuild:
             if reset_dset is not None:
                 print("Resetting Dataloader")
                 print("Putting in params from previous dataset.")
@@ -145,13 +304,8 @@ def reconstruct(
                     self.scheduler_params = scheduler_params
                     self.set_schedulers(self.scheduler_params, num_iter=num_iter)
 
-            self.dataloader, self.sampler, self.val_dataloader, self.val_sampler = (
-                self.setup_dataloader(
-                    self.dset,
-                    batch_size,
-                    num_workers=num_workers,
-                    val_fraction=val_fraction,
-                )
+            self._setup_recon_dataloaders(
+                batch_size, num_workers, val_fraction, holdout_fraction, holdout_seed
             )
 
         # Type check for INR-based reconstruction
@@ -160,6 +314,41 @@ def reconstruct(
                 "Only TomographyINRDataset is supported for this reconstruction method."
             )
 
+        pose_warmup_pending = pose_warmup_epochs > 0 and self.dset.has_optimizer()
+        if pose_warmup_pending:
+            # remove_optimizer also clears the stored specs; restore those specs without
+            # rebuilding the optimizer/scheduler until the warmup finishes.
+            pose_optimizer_params = self.optimizer_params["pose"]
+            pose_scheduler_params = self.scheduler_params["pose"]
+            self.dset.remove_optimizer()
+            self.dset.optimizer_params = pose_optimizer_params
+            self.dset.scheduler_params = pose_scheduler_params
+
+        capture_enabled = cuda_graphs
+        if capture_enabled and (self.dset.has_optimizer() or pose_warmup_pending):
+            warnings.warn(
+                "cuda_graphs=True is falling back to eager reconstruction because pose "
+                "learning is enabled.",
+                stacklevel=2,
+            )
+            capture_enabled = False
+
+        tilted_model = (
+            self.obj_model.model if isinstance(self.obj_model.model, KPlanesTILTED) else None
+        )
+        static_rotation_matrices: torch.Tensor | None = None
+        if capture_enabled and tilted_model is not None:
+            if any(parameter.requires_grad for parameter in tilted_model.so3.parameters()):
+                warnings.warn(
+                    "cuda_graphs=True is falling back to eager reconstruction because "
+                    "KPlanesTILTED SO(3) learning is enabled.",
+                    stacklevel=2,
+                )
+                capture_enabled = False
+            else:
+                with torch.no_grad():
+                    static_rotation_matrices = tilted_model.so3.as_matrix().detach()
+
         N = max(self.obj_model.shape)
 
         if num_samples_per_ray is None:
@@ -177,8 +366,138 @@ def reconstruct(
 
         loss_func = get_loss_module(name=loss_type, dtype=self.obj_model.dtype, **loss_func_kwargs)
 
+        forward_graph: torch.cuda.CUDAGraph | None = None
+        optimizer_graph: torch.cuda.CUDAGraph | None = None
+        graph_pool = None
+        static_batch: dict[str, torch.Tensor] | None = None
+        static_losses: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None
+        graph_grads: list[tuple[torch.nn.Parameter, torch.Tensor]] = []
+        graph_num_samples_per_ray: int | None = None
+        optimizer = self.obj_model.optimizer
+        optimizer_graph_eligible = (
+            capture_enabled
+            and not grad_scaler_enabled
+            and isinstance(optimizer, torch.optim.Adam)
+            and all(
+                group.get("capturable", False) and group.get("fused", False)
+                for group in optimizer.param_groups
+            )
+        )
+        pred_fork_enabled = (
+            self.device.type == "cuda"
+            and not capture_enabled
+            and getattr(self.obj_model.constraints, "s3im_weight", 0.0) > 0
+            # The side-stream schedule assumes RNG-free, stateless consistency
+            # losses. Keep it disabled for user-supplied stochastic losses.
+            and os.environ.get("QUANTEM_RECON_PRED_FORK", "0") != "0"
+        )
+        loss_branch_stream = torch.cuda.Stream(device=self.device) if pred_fork_enabled else None
+        pred_ready_event = torch.cuda.Event() if pred_fork_enabled else None
+        consistency_ready_event = torch.cuda.Event() if pred_fork_enabled else None
+
+        def optimizer_state_is_initialized() -> bool:
+            assert optimizer is not None
+            return all(
+                parameter.grad is None or bool(optimizer.state.get(parameter))
+                for group in optimizer.param_groups
+                for parameter in group["params"]
+            )
+
+        def clear_model_grads(*, set_to_none: bool) -> None:
+            for parameter in self.obj_model.model.parameters():
+                if set_to_none:
+                    parameter.grad = None
+                elif parameter.grad is not None:
+                    parameter.grad.zero_()
+
+        def forward_loss_backward(
+            batch: dict[str, torch.Tensor],
+            samples_per_ray: int,
+        ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+            with torch.autocast(
+                device_type=self.device.type,
+                dtype=torch_autocast_dtype,
+                enabled=autocast_enabled,
+            ):
+                nvtx.range_push("get_coords")
+                all_coords = self.dset.get_coords(batch, N, samples_per_ray)
+                nvtx.range_pop()
+
+                nvtx.range_push("obj_forward")
+                tap_coords = self.obj_model.sample_tv_tap_coords(all_coords)
+                with self.obj_model.reconstruction_forward_context():
+                    if tap_coords is not None:
+                        all_densities, tv_tap_raw = self.obj_model.forward_with_tv_taps(
+                            all_coords, tap_coords
+                        )
+                    else:
+                        all_densities = self.obj_model.forward(all_coords)
+                        tv_tap_raw = None
+                nvtx.range_pop()
+
+                nvtx.range_push("integrate_rays")
+                integrated_densities = self.dset.integrate_rays(
+                    all_densities,
+                    samples_per_ray,
+                    len(batch["target_value"]),
+                )
+                nvtx.range_pop()
+
+            pred = integrated_densities.float()
+            target = batch["target_value"].to(self.device, non_blocking=True).float()
+
+            if loss_branch_stream is not None:
+                assert pred_ready_event is not None
+                assert consistency_ready_event is not None
+                main_stream = torch.cuda.current_stream(self.device)
+                pred_ready_event.record(main_stream)
+                loss_branch_stream.wait_event(pred_ready_event)
+                pred.record_stream(loss_branch_stream)
+                target.record_stream(loss_branch_stream)
+                nvtx.range_push("consistency_loss")
+                with torch.cuda.stream(loss_branch_stream):
+                    batch_consistency_loss = loss_func(pred, target)
+                    consistency_ready_event.record(loss_branch_stream)
+                nvtx.range_pop()
+
+            nvtx.range_push("soft_constraints")
+            constraint_densities = self.dset.graph_constraint_densities(all_densities)
+            soft_constraints_loss = self.obj_model.apply_soft_constraints(
+                ctx=ReconstructionContext(
+                    coords=all_coords,
+                    pred=pred,
+                    all_densities=constraint_densities,
+                    target=target,
+                    tv_tap_densities=tv_tap_raw,
+                )
+            )
+            soft_constraints_loss += self.dset.apply_soft_constraints()
+            nvtx.range_pop()
+
+            if loss_branch_stream is None:
+                nvtx.range_push("consistency_loss")
+                batch_consistency_loss = loss_func(pred, target)
+                nvtx.range_pop()
+            else:
+                main_stream.wait_event(consistency_ready_event)
+                batch_consistency_loss.record_stream(main_stream)
+            batch_loss = batch_consistency_loss.float() + soft_constraints_loss.float()
+
+            nvtx.range_push("backward")
+            if grad_scaler_enabled:
+                grad_scaler.scale(batch_loss).backward()
+            else:
+                batch_loss.backward()
+            nvtx.range_pop()
+            return batch_loss, batch_consistency_loss, soft_constraints_loss
+
         pbar = tqdm(range(num_iter), disable=not self.verbose)
         for a0 in pbar:
+            if pose_warmup_pending and a0 == pose_warmup_epochs:
+                self.dset.set_optimizer(self.optimizer_params["pose"])
+                self.dset.set_scheduler(self.scheduler_params["pose"], num_iter=num_iter)
+                pose_warmup_pending = False
+            nvtx.range_push(f"epoch_{a0}")
             consistency_loss = torch.tensor(0.0, device=self.device)
             total_loss = torch.tensor(0.0, device=self.device)
             epoch_soft_constraint_loss = torch.tensor(0.0, device=self.device)
@@ -202,48 +521,149 @@ def reconstruct(
                 curr_num_samples_per_ray = num_samples_per_ray
 
             for batch_idx, batch in enumerate(self.dataloader):
-                self.zero_grad_all()
-                with torch.autocast(
-                    device_type=self.device.type,
-                    dtype=torch.bfloat16,
-                    enabled=False,
-                ):
-                    all_coords = self.dset.get_coords(batch, N, curr_num_samples_per_ray)
-
-                    all_densities = self.obj_model.forward(all_coords)
-
-                    integrated_densities = self.dset.integrate_rays(
-                        all_densities,
-                        curr_num_samples_per_ray,
-                        len(batch["target_value"]),
-                    )
-
-                pred = integrated_densities.float()
-
-                soft_constraints_loss = self.obj_model.apply_soft_constraints(
-                    ctx=ReconstructionContext(
-                        coords=all_coords,
-                        pred=pred,
-                        all_densities=all_densities,
+                nvtx.range_push("batch")
+                full_batch = len(batch["target_value"]) == batch_size
+                graph_shape_matches = (
+                    forward_graph is not None
+                    and curr_num_samples_per_ray == graph_num_samples_per_ray
+                    and static_batch is not None
+                    and all(
+                        key in batch
+                        and batch[key].shape == value.shape
+                        and batch[key].dtype == value.dtype
+                        for key, value in static_batch.items()
                     )
                 )
-
-                target = batch["target_value"].to(self.device, non_blocking=True).float()
-
-                batch_consistency_loss = loss_func(pred, target)
-
-                soft_constraints_loss += self.dset.apply_soft_constraints()
+                used_forward_graph = False
+
+                if capture_enabled and full_batch and forward_graph is None:
+                    static_batch = {
+                        key: value.detach().to(self.device).clone() for key, value in batch.items()
+                    }
+                    graph_num_samples_per_ray = curr_num_samples_per_ray
+                    try:
+                        self.dset._cuda_graph_static_sampling = True
+                        if tilted_model is not None:
+                            assert static_rotation_matrices is not None
+                            tilted_model._rotation_matrices_override = static_rotation_matrices
+
+                        # Materialize lazy kernels/optimizer scale state on a side stream.
+                        current_stream = torch.cuda.current_stream(self.device)
+                        warmup_stream = torch.cuda.Stream(device=self.device)
+                        warmup_stream.wait_stream(current_stream)
+                        with torch.cuda.stream(warmup_stream):
+                            for _ in range(3):
+                                clear_model_grads(set_to_none=True)
+                                forward_loss_backward(
+                                    static_batch,
+                                    curr_num_samples_per_ray,
+                                )
+                        current_stream.wait_stream(warmup_stream)
+
+                        graph_pool = torch.cuda.graph_pool_handle()
+                        forward_graph = torch.cuda.CUDAGraph()
+                        with torch.cuda.graph(forward_graph, pool=graph_pool):
+                            clear_model_grads(set_to_none=False)
+                            static_losses = forward_loss_backward(
+                                static_batch,
+                                curr_num_samples_per_ray,
+                            )
+                    finally:
+                        self.dset._cuda_graph_static_sampling = False
+                        if tilted_model is not None:
+                            tilted_model._rotation_matrices_override = None
+                    graph_grads = [
+                        (parameter, parameter.grad)
+                        for parameter in self.obj_model.model.parameters()
+                        if parameter.grad is not None
+                    ]
+                    # Capture records the work but does not populate static outputs for
+                    # this first batch; replay once before consuming its losses/gradients.
+                    forward_graph.replay()
+                    batch_loss, batch_consistency_loss, soft_constraints_loss = static_losses
+                    used_forward_graph = True
+                elif capture_enabled and full_batch and graph_shape_matches:
+                    assert static_batch is not None
+                    assert static_losses is not None
+                    for key, value in static_batch.items():
+                        value.copy_(batch[key], non_blocking=True)
+                    for parameter, static_grad in graph_grads:
+                        parameter.grad = static_grad
+                    forward_graph.replay()
+                    batch_loss, batch_consistency_loss, soft_constraints_loss = static_losses
+                    used_forward_graph = True
+                else:
+                    # DeviceBatchSampler drops the tail. This path protects custom
+                    # samplers/wrappers and sample-count schedules without violating a
+                    # previously captured graph's static-shape contract.
+                    self.zero_grad_all()
+                    batch_loss, batch_consistency_loss, soft_constraints_loss = (
+                        forward_loss_backward(batch, curr_num_samples_per_ray)
+                    )
 
                 epoch_soft_constraint_loss += soft_constraints_loss.detach()
-
-                batch_loss = batch_consistency_loss.float() + soft_constraints_loss.float()
-
-                batch_loss.backward()
                 # Clip gradients
-                torch.nn.utils.clip_grad_norm_(self.obj_model.model.parameters(), max_norm=1.0)
-                self.step_optimizers()
+                nvtx.range_push("clip_and_optim_step")
+                self._sync_pose_gradients_ddp()
+                if grad_scaler_enabled:
+                    if "object" in self.optimizer_params and self.obj_model.has_optimizer():
+                        grad_scaler.unscale_(self.obj_model.optimizer)
+                    if "pose" in self.optimizer_params and self.dset.has_optimizer():
+                        grad_scaler.unscale_(self.dset.optimizer)
+                if grad_clip_max_norm is not None and grad_clip_max_norm > 0:
+                    torch.nn.utils.clip_grad_norm_(
+                        self.obj_model.model.parameters(), max_norm=grad_clip_max_norm
+                    )
+                if grad_scaler_enabled:
+                    scaler_stepped = False
+                    if "object" in self.optimizer_params and self.obj_model.has_optimizer():
+                        grad_scaler.step(self.obj_model.optimizer)
+                        scaler_stepped = True
+                    if "pose" in self.optimizer_params and self.dset.has_optimizer():
+                        grad_scaler.step(self.dset.optimizer)
+                        scaler_stepped = True
+                    if scaler_stepped:
+                        grad_scaler.update()
+                elif optimizer_graph_eligible and graph_pool is not None and used_forward_graph:
+                    assert optimizer is not None
+                    if not optimizer_state_is_initialized():
+                        # Adam lazily creates and zeroes its state on the first step. If
+                        # that initialization is captured, every replay resets ``step``
+                        # before incrementing it, freezing bias correction at step 1.
+                        # Execute the first logical update eagerly so subsequent capture
+                        # sees stable state tensors.
+                        optimizer.step()
+                    elif optimizer_graph is None:
+                        optimizer_graph = torch.cuda.CUDAGraph()
+                        with torch.cuda.graph(optimizer_graph, pool=graph_pool):
+                            optimizer.step()
+                        # Capture records the step but does not execute this iteration's
+                        # update. Mirror the forward graph and replay it immediately.
+                        optimizer_graph.replay()
+                    else:
+                        optimizer_graph.replay()
+                else:
+                    self.step_optimizers()
+                nvtx.range_pop()
+                graph_step_callback = getattr(self, "_cuda_graph_step_callback", None)
+                if capture_enabled and graph_step_callback is not None:
+                    graph_step_callback()
+                self._grad_steps = getattr(self, "_grad_steps", 0) + 1
+                nsys_capture_tick(self._grad_steps)
+                if snapshots_enabled and _should_take_grad_step_snapshot(
+                    self._grad_steps, snapshot_every
+                ):
+                    _take_grad_step_snapshot(
+                        obj_model=self.obj_model,
+                        grad_step=self._grad_steps,
+                        global_rank=self.global_rank,
+                        logger=self.logger,
+                        snapshot_dir=snapshot_dir,
+                        snapshot_callback=snapshot_callback,
+                    )
                 total_loss += batch_loss.detach()
                 consistency_loss += batch_consistency_loss.detach()
+                nvtx.range_pop()  # batch
 
             if isinstance(self.obj_model.model, CPTilted):
                 if a0 == 0:
@@ -264,63 +684,59 @@ def reconstruct(
                     )
                     prev_R = R_now.clone()
 
+            # One stacked all_reduce and one host sync instead of three of each.
+            losses = torch.stack([total_loss, consistency_loss, epoch_soft_constraint_loss])
             if self.world_size > 1:
-                dist.all_reduce(total_loss, dist.ReduceOp.AVG)
-                dist.all_reduce(consistency_loss, dist.ReduceOp.AVG)
-                dist.all_reduce(epoch_soft_constraint_loss, dist.ReduceOp.AVG)
-
-            total_loss = total_loss.item() / len(self.dataloader)
-            consistency_loss = consistency_loss.item() / len(self.dataloader)
-            epoch_soft_constraint_loss = epoch_soft_constraint_loss.item() / len(self.dataloader)
+                nvtx.range_push("ddp_allreduce_epoch_metrics")
+                dist.all_reduce(losses, dist.ReduceOp.AVG)
+                nvtx.range_pop()
+            total_loss, consistency_loss, epoch_soft_constraint_loss = (
+                losses / len(self.dataloader)
+            ).tolist()
 
             self.step_schedulers(loss=total_loss)
             # TODO: Maybe reorganize the losses so that the order makes sense lol.
 
             avg_val_loss = None
-            if self.val_dataloader is not None:
-                print("Validating...")
-                self.obj_model.model.eval()
-                self.dset.eval()
-                with torch.no_grad():
-                    val_loss = torch.tensor(0.0, device=self.device)
-
-                    for batch in self.val_dataloader:
-                        with torch.autocast(
-                            device_type=self.device.type,
-                            dtype=torch.bfloat16,
-                            enabled=True,
-                        ):
-                            all_coords = self.dset.get_coords(batch, N, curr_num_samples_per_ray)
-
-                            all_densities = self.obj_model.forward(all_coords)
-
-                            integrated_densities = self.dset.integrate_rays(
-                                all_densities,
-                                curr_num_samples_per_ray,
-                                len(batch["target_value"]),
-                            )
-
-                            target = (
-                                batch["target_value"].to(self.device, non_blocking=True).float()
-                            )
-
-                            batch_val_loss = torch.nn.functional.mse_loss(
-                                integrated_densities, target
-                            )
-
-                            val_loss += batch_val_loss.detach()
-
-                    avg_val_loss = val_loss.item() / len(self.val_dataloader)
-
-            metrics = torch.tensor(
-                [total_loss, consistency_loss, epoch_soft_constraint_loss], device=self.device
+            avg_val_fg_loss = None
+            avg_val_bg_loss = None
+            validate_this_epoch = self.val_dataloader is not None and (
+                holdout_fraction <= 0.0 or (a0 + 1) % holdout_every == 0
             )
+            if validate_this_epoch:
+                print("Validating...")
+                nvtx.range_push("validation")
+                avg_val_loss = self._evaluate_validation_loss(
+                    dataloader=self.val_dataloader,
+                    num_samples_per_ray=curr_num_samples_per_ray,
+                    object_extent=N,
+                    loss_func=loss_func,
+                    autocast_dtype=torch_autocast_dtype,
+                    autocast_enabled=autocast_enabled,
+                )
+                if getattr(self, "val_fg_dataloader", None) is not None:
+                    avg_val_fg_loss = self._evaluate_validation_loss(
+                        dataloader=self.val_fg_dataloader,
+                        num_samples_per_ray=curr_num_samples_per_ray,
+                        object_extent=N,
+                        loss_func=loss_func,
+                        autocast_dtype=torch_autocast_dtype,
+                        autocast_enabled=autocast_enabled,
+                    )
+                if getattr(self, "val_bg_dataloader", None) is not None:
+                    avg_val_bg_loss = self._evaluate_validation_loss(
+                        dataloader=self.val_bg_dataloader,
+                        num_samples_per_ray=curr_num_samples_per_ray,
+                        object_extent=N,
+                        loss_func=loss_func,
+                        autocast_dtype=torch_autocast_dtype,
+                        autocast_enabled=autocast_enabled,
+                    )
+                nvtx.range_pop()  # validation
 
-            if self.world_size > 1:
-                dist.all_reduce(metrics, dist.ReduceOp.AVG)
-
-            total_loss, consistency_loss, epoch_soft_constraint_loss = metrics.tolist()
-
+            # The three losses were already rank-averaged (and batch-normalized) right
+            # after the batch loop; re-reducing identical values here was a redundant
+            # all_reduce plus an extra host sync per epoch.
             pbar.set_description(
                 f"Reconstruction | Loss: {total_loss:.5e}, Consistency Loss: {consistency_loss:.5e}, Soft Constraint Loss: {epoch_soft_constraint_loss:.5e}"
             )
@@ -357,7 +773,10 @@ def reconstruct(
                         total_loss=total_loss,
                         learning_rates=self.get_current_lrs(),
                         num_samples_per_ray=curr_num_samples_per_ray,
-                        val_loss=avg_val_loss if self.val_dataloader is not None else None,
+                        val_loss=avg_val_loss if validate_this_epoch else None,
+                        val_fg_loss=avg_val_fg_loss,
+                        val_bg_loss=avg_val_bg_loss,
+                        grad_step=self._grad_steps,
                     )
 
                 self.logger.flush()
@@ -366,11 +785,69 @@ def reconstruct(
                     print(
                         f"Reconstruction Epoch {self.num_epochs} | Loss: {total_loss:.5e}, Consistency Loss: {consistency_loss:.5e}, Soft Constraint Loss: {epoch_soft_constraint_loss:.5e}"
                     )
+
+            # Opt-in per-epoch evaluation hook (e.g. SSIM-vs-ground-truth convergence
+            # curves). Note `obj_view` is collective, so the callback must be invoked on
+            # every rank; it is the caller's responsibility to guard rank-0-only work.
+            if eval_callback is not None and eval_every > 0 and (a0 + 1) % eval_every == 0:
+                eval_callback(a0 + 1)
+            nvtx.range_pop()  # epoch
+
         if show_metrics and self.world_size == 1:
             self.plot_losses()
 
     # --- Helper Functions ---
 
+    def _evaluate_validation_loss(
+        self,
+        *,
+        dataloader: DeviceBatchSampler,
+        num_samples_per_ray: int,
+        object_extent: int,
+        loss_func: torch.nn.Module,
+        autocast_dtype: torch.dtype = torch.bfloat16,
+        autocast_enabled: bool = False,
+    ) -> float | None:
+        val_loss = torch.tensor(0.0, device=self.device)
+        val_batches = torch.tensor(0.0, device=self.device)
+        model_was_training = self.obj_model.model.training
+        dset_was_training = self.dset.training
+
+        self.obj_model.model.eval()
+        self.dset.eval()
+        try:
+            with torch.no_grad():
+                for batch in dataloader:
+                    # Keep validation autocast consistent with the training pass.
+                    with torch.autocast(
+                        device_type=self.device.type,
+                        dtype=autocast_dtype,
+                        enabled=autocast_enabled,
+                    ):
+                        all_coords = self.dset.get_coords(
+                            batch, object_extent, num_samples_per_ray
+                        )
+                        all_densities = self.obj_model.forward(all_coords)
+                        integrated_densities = self.dset.integrate_rays(
+                            all_densities,
+                            num_samples_per_ray,
+                            len(batch["target_value"]),
+                        )
+                        target = batch["target_value"].to(self.device, non_blocking=True).float()
+                        batch_val_loss = loss_func(integrated_densities.float(), target)
+                        val_loss += batch_val_loss.detach()
+                        val_batches += 1.0
+        finally:
+            self.obj_model.model.train(model_was_training)
+            self.dset.train(dset_was_training)
+
+        stats = torch.stack([val_loss, val_batches])
+        if self.world_size > 1:
+            dist.all_reduce(stats, dist.ReduceOp.SUM)
+        if stats[1].item() == 0.0:
+            return None
+        return (stats[0] / stats[1]).item()
+
     def save_volume(self, path: str = "recon_volume.npz", overwrite: bool = False):
         """
         Saves volume to a numpy array file. Does not save the full Tomography object.
@@ -404,13 +881,124 @@ def from_file(
             batch_size=tomography.batch_size,
             num_workers=tomography.num_workers,
             val_fraction=tomography.val_fraction,
+            holdout_fraction=getattr(tomography, "holdout_fraction", 0.0),
+            holdout_seed=getattr(tomography, "holdout_seed", 0),
         )
         return tomography
 
-    def _rebuild_dataloader(self, batch_size: int, num_workers: int, val_fraction: float):
+    def _rebuild_dataloader(
+        self,
+        batch_size: int,
+        num_workers: int,
+        val_fraction: float,
+        holdout_fraction: float = 0.0,
+        holdout_seed: int = 0,
+    ):
         """
         Rebuilds the dataloader due to persistent workers error when reloading the object.
         """
+        self._setup_recon_dataloaders(
+            batch_size, num_workers, val_fraction, holdout_fraction, holdout_seed
+        )
+
+    def _setup_recon_dataloaders(
+        self,
+        batch_size: int,
+        num_workers: int,
+        val_fraction: float,
+        holdout_fraction: float = 0.0,
+        holdout_seed: int = 0,
+    ):
+        """Build the train/val batch iterators.
+
+        INR datasets use ``DeviceBatchSampler`` — batches are built with
+        tensor ops on the compute device from a device-resident tilt stack,
+        removing the per-pixel ``__getitem__`` / collate / H2D-copy
+        dataloader bottleneck (``num_workers`` is ignored on this path).
+        Distributed runs shard the same seeded epoch permutation across
+        ranks (DistributedSampler semantics; the loop's ``set_epoch`` drives
+        reshuffling). Non-INR datasets keep the DataLoader path.
+        """
+        self.val_fg_dataloader = None
+        self.val_bg_dataloader = None
+        if isinstance(self.dset, TomographyINRDataset):
+            n = len(self.dset)
+            if holdout_fraction > 0.0:
+                split = build_pixel_holdout_split(
+                    self.dset.tilt_stack,
+                    holdout_fraction=holdout_fraction,
+                    holdout_seed=holdout_seed,
+                )
+                train_indices = split.train_indices
+                val_indices = split.val_indices
+                val_fg_indices = split.val_fg_indices
+                val_bg_indices = split.val_bg_indices
+            else:
+                n_val = int(n * val_fraction)
+                # Fixed-seed split: identical across DDP ranks (no train/val
+                # leakage between ranks) and stable across save/reload, so a
+                # resumed run keeps validating on the same held-out pixels.
+                split_gen = torch.Generator()
+                split_gen.manual_seed(0)
+                perm = torch.randperm(n, generator=split_gen)
+                train_indices = perm[n_val:]
+                val_indices = perm[:n_val]
+                val_fg_indices = torch.empty(0, dtype=torch.long)
+                val_bg_indices = torch.empty(0, dtype=torch.long)
+
+            ddp = dict(rank=self.global_rank, world_size=self.world_size)
+            self.dataloader = DeviceBatchSampler(
+                self.dset, batch_size, self.device, indices=train_indices, **ddp
+            )
+            # The val sampler keeps its own device-resident copy of the tilt
+            # stack; acceptable, since val_fraction > 0 is the rare case.
+            val = (
+                DeviceBatchSampler(
+                    self.dset,
+                    batch_size,
+                    self.device,
+                    indices=val_indices,
+                    shuffle=False,
+                    drop_last=False,
+                    **ddp,
+                )
+                if len(val_indices) > 0
+                else None
+            )
+            self.val_dataloader = val if val is not None and len(val) > 0 else None
+            val_fg = (
+                DeviceBatchSampler(
+                    self.dset,
+                    batch_size,
+                    self.device,
+                    indices=val_fg_indices,
+                    shuffle=False,
+                    drop_last=False,
+                    **ddp,
+                )
+                if len(val_fg_indices) > 0
+                else None
+            )
+            val_bg = (
+                DeviceBatchSampler(
+                    self.dset,
+                    batch_size,
+                    self.device,
+                    indices=val_bg_indices,
+                    shuffle=False,
+                    drop_last=False,
+                    **ddp,
+                )
+                if len(val_bg_indices) > 0
+                else None
+            )
+            self.val_fg_dataloader = val_fg if val_fg is not None and len(val_fg) > 0 else None
+            self.val_bg_dataloader = val_bg if val_bg is not None and len(val_bg) > 0 else None
+            # The training loop calls set_epoch on self.sampler.
+            self.sampler = self.dataloader
+            self.val_sampler = None
+            return
+
         self.dataloader, self.sampler, self.val_dataloader, self.val_sampler = (
             self.setup_dataloader(
                 self.dset,
@@ -566,7 +1154,10 @@ def _reconstruction_epoch(
         if inline_alignment:
             for ind in range(len(self.dset.tilt_angles)):
                 im_proj = proj_forward[:, ind, :]
-                im_meas = self.dset.forward(ind).target  # type: ignore
+                # proj_forward rows are volume slices (the tilt axis), i.e. the
+                # transpose of the stored tilt image -- the same orientation the
+                # error term below compares against.
+                im_meas = self.dset.forward(ind).target.T  # type: ignore
                 shift = torch_phase_cross_correlation(im_proj, im_meas)
                 if torch.linalg.norm(shift) <= 32:
                     shifted = torch.fft.ifft2(
@@ -585,7 +1176,11 @@ def _reconstruction_epoch(
                         )
                     ).real
 
-                    proj_forward[:, ind, :] = shifted
+                    # Persist the aligned measurement in the tilt stack: the error
+                    # term reads the stack, and proj_forward is overwritten by
+                    # radon_torch below, so writing the aligned image there
+                    # silently discarded the alignment.
+                    self.dset.tilt_stack[ind] = shifted.T
 
         if mode == "sirt" or mode == "fbp":
             proj_forward = radon_torch(
diff --git a/src/quantem/tomography/tomography_base.py b/src/quantem/tomography/tomography_base.py
index 419b5ede6..ee2857b78 100644
--- a/src/quantem/tomography/tomography_base.py
+++ b/src/quantem/tomography/tomography_base.py
@@ -51,6 +51,7 @@ def __init__(
         self._consistency_losses: list[float] = []
         self._val_losses: list[float] = []
         self._lrs: dict[str, list] = {}
+        self._grad_steps: int = 0
         # DDP Initialization
         if isinstance(obj_model, ObjectINR) or isinstance(obj_model, ObjectTensorDecomp):
             self.setup_distributed(device=device)
@@ -58,7 +59,11 @@ def __init__(
                 print("Setting up DDP for obj_model")
 
         self.dset = dset
-        self.dset.to(device)
+        # Use self.device (set by setup_distributed to cuda:local_rank under DDP),
+        # NOT the local `device` arg, which stays "cuda:0" on every rank and would
+        # strand the dataset's pose params on cuda:0 while the model/sampler live
+        # on cuda:local_rank (device-mismatch in _apply_pose's batched matmul).
+        self.dset.to(self.device)
 
     # --- Properties ---
     @property
diff --git a/src/quantem/tomography/tomography_context.py b/src/quantem/tomography/tomography_context.py
index d322287c6..b62078e22 100644
--- a/src/quantem/tomography/tomography_context.py
+++ b/src/quantem/tomography/tomography_context.py
@@ -1,9 +1,10 @@
 from dataclasses import dataclass
 from typing import Optional
-from quantem.core.ml.constraints import BaseContext
 
 import torch
 
+from quantem.core.ml.constraints import BaseContext
+
 
 @dataclass
 class ReconstructionContext(BaseContext):
@@ -21,6 +22,7 @@ class ReconstructionContext(BaseContext):
     - pred: Predicted values per coordinate position from the model.
     - all_densities: Integrated densities per ray from the model.
     - obj: Object model (INR, TensorDecomp, etc.).
+    - target: Measured projection pixels for the batch, for data-coupled constraints (e.g. S3IM).
     """
 
     volume: Optional[torch.Tensor] = None
@@ -28,3 +30,5 @@ class ReconstructionContext(BaseContext):
     pred: Optional[torch.Tensor] = None
     all_densities: Optional[torch.Tensor] = None
     obj: Optional[torch.Tensor] = None
+    target: Optional[torch.Tensor] = None
+    tv_tap_densities: Optional[torch.Tensor] = None
diff --git a/src/quantem/tomography/tomography_lite.py b/src/quantem/tomography/tomography_lite.py
index 7bd9ad674..48ccf3edd 100644
--- a/src/quantem/tomography/tomography_lite.py
+++ b/src/quantem/tomography/tomography_lite.py
@@ -1,5 +1,6 @@
 import os
-from typing import Literal, Self
+from pathlib import Path
+from typing import Callable, Literal, Self
 
 import numpy as np
 import torch
@@ -94,6 +95,12 @@ def reconstruct(  # type:ignore[reportIncompatibleMethodOverride] ## easier than
         obj_constraints: ObjConstraintsType | dict | None = None,
         dset_constraints: DatasetConstraintsType | dict | None = None,
         show_metrics: bool = False,
+        holdout_fraction: float = 0.0,
+        holdout_seed: int = 0,
+        holdout_every: int = 1,
+        snapshot_every: int = 0,
+        snapshot_dir: str | Path | None = None,
+        snapshot_callback: Callable[[int, np.ndarray | None], None] | None = None,
     ):
         if self.num_epochs == 0:
             opt_params = {
@@ -135,6 +142,12 @@ def reconstruct(  # type:ignore[reportIncompatibleMethodOverride] ## easier than
             obj_constraints=obj_constraints,
             dset_constraints=dset_constraints,
             show_metrics=show_metrics,
+            holdout_fraction=holdout_fraction,
+            holdout_seed=holdout_seed,
+            holdout_every=holdout_every,
+            snapshot_every=snapshot_every,
+            snapshot_dir=snapshot_dir,
+            snapshot_callback=snapshot_callback,
         )
 
 
diff --git a/src/quantem/tomography/tomography_opt.py b/src/quantem/tomography/tomography_opt.py
index 7f98fabd1..a4a824f5a 100644
--- a/src/quantem/tomography/tomography_opt.py
+++ b/src/quantem/tomography/tomography_opt.py
@@ -7,9 +7,29 @@
     OptimizerParamsType,
     SchedulerParamsType,
 )
+from quantem.tomography.dataset_models import TomographyINRDataset
 from quantem.tomography.tomography_base import TomographyBase
 
 
+def _pose_optimizer_specs(
+    dset: TomographyINRDataset,
+    pose_lr: float,
+    shift_lr: float | None = None,
+    tilt_axis_lr: float | None = None,
+) -> dict[str, OptimizerParams.Adam]:
+    """Build role-keyed Adam specs, falling back to the legacy shared pose LR."""
+    specs = {}
+    if dset.learn_shift:
+        specs["pose_shift"] = OptimizerParams.Adam(
+            lr=pose_lr if shift_lr is None else shift_lr
+        )
+    if dset.learn_tilt_axis:
+        specs["pose_tilt_axis"] = OptimizerParams.Adam(
+            lr=pose_lr if tilt_axis_lr is None else tilt_axis_lr
+        )
+    return specs
+
+
 class TomographyOpt(TomographyBase):
     """
     Class for handling all the optimizers and schedulers for the tomography reconstruction.
diff --git a/src/quantem/tomography/utils.py b/src/quantem/tomography/utils.py
index 08093925a..553a27b46 100644
--- a/src/quantem/tomography/utils.py
+++ b/src/quantem/tomography/utils.py
@@ -1,14 +1,19 @@
 import torch
 import torch.nn.functional as F
 
+from quantem.core import config
+
 # --- Projection Operator Utils ---
 
 
 def rot_ZXZ(mags, z1, x, z3, device, mode="bilinear"):
-    if not isinstance(x, torch.Tensor) or not isinstance(z1, torch.Tensor):
-        z1 = torch.tensor(z1, dtype=torch.float32, device=device)
-        x = torch.tensor(x, dtype=torch.float32, device=device)
-        z3 = torch.tensor(z3, dtype=torch.float32, device=device)
+    # Convert each angle independently: re-wrapping an existing tensor with
+    # torch.tensor() copies and detaches it, silently cutting gradient flow
+    # through tensor angles whenever any other angle is passed as a float.
+    z1, x, z3 = (
+        a if isinstance(a, torch.Tensor) else torch.tensor(a, dtype=torch.float32, device=device)
+        for a in (z1, x, z3)
+    )
     curr_mags = mags
 
     curr_mags = differentiable_rotz_vectorized(curr_mags, z1, mode)
@@ -19,54 +24,90 @@ def rot_ZXZ(mags, z1, x, z3, device, mode="bilinear"):
     return curr_mags
 
 
-def differentiable_rotz_vectorized(mags, theta, mode="bilinear"):
-    _, dimz, dimy, dimx = mags.shape
-
-    if theta.dim() == 0:
-        theta = theta.unsqueeze(0)
-
+def _rot2d_affine_matrix(theta: torch.Tensor, batch: int) -> torch.Tensor:
+    """(B, 2, 3) in-plane rotation matrices, broadcasting a single angle over the batch."""
     theta_rad = torch.deg2rad(theta)
-
     cos_t, sin_t = torch.cos(theta_rad), torch.sin(theta_rad)
     affine_matrix = torch.stack(
         [cos_t, -sin_t, torch.zeros_like(theta), sin_t, cos_t, torch.zeros_like(theta)], dim=-1
     ).view(-1, 2, 3)
+    if affine_matrix.shape[0] == 1 and batch > 1:
+        affine_matrix = affine_matrix.expand(batch, 2, 3)
+    elif affine_matrix.shape[0] != batch and batch != 1:
+        raise ValueError(
+            f"Got {affine_matrix.shape[0]} angles for a batch of {batch} volumes; "
+            "pass one angle, or one angle per volume."
+        )
+    return affine_matrix
 
-    mags = mags.permute(1, 0, 2, 3)
 
-    def transform_slice(mag_slice):
-        grid = F.affine_grid(affine_matrix, mag_slice.unsqueeze(0).shape, align_corners=False)
-        return F.grid_sample(mag_slice.unsqueeze(0), grid, mode=mode, align_corners=False).squeeze(
-            0
-        )
+def differentiable_rotz_vectorized(mags, theta, mode="bilinear"):
+    B, dimz, dimy, dimx = mags.shape
+
+    if theta.dim() == 0:
+        theta = theta.unsqueeze(0)
 
-    rotated_mags = torch.vmap(transform_slice)(mags)
-    return rotated_mags.permute(1, 0, 2, 3)
+    # A rotation about z applies the same 2-D transform to every z-slice, so the
+    # z-axis rides along as grid_sample channels in a single call. (The previous
+    # per-slice vmap also broke for more than one angle, because affine_grid
+    # requires the matrix batch to match the slice batch of 1.)
+    affine_matrix = _rot2d_affine_matrix(theta, B)
+    if affine_matrix.shape[0] > B:  # one volume, many angles
+        mags = mags.expand(affine_matrix.shape[0], dimz, dimy, dimx)
+    grid = F.affine_grid(affine_matrix, mags.shape, align_corners=False)
+    return F.grid_sample(mags, grid, mode=mode, align_corners=False)
 
 
 def differentiable_rotx_vectorized(mags, theta, mode="bilinear"):
-    _, dimz, dimy, dimx = mags.shape
+    B, dimz, dimy, dimx = mags.shape
 
     if theta.dim() == 0:
         theta = theta.unsqueeze(0)
 
-    theta_rad = torch.deg2rad(theta)
+    # Same trick as rotz with the x-axis as the channel dim: rotate in (z, y).
+    affine_matrix = _rot2d_affine_matrix(theta, B)
+    mags = mags.permute(0, 3, 1, 2)  # (B, X, Z, Y)
+    if affine_matrix.shape[0] > B:  # one volume, many angles
+        mags = mags.expand(affine_matrix.shape[0], dimx, dimz, dimy)
+    grid = F.affine_grid(affine_matrix, mags.shape, align_corners=False)
+    rotated = F.grid_sample(mags, grid, mode=mode, align_corners=False)
+    return rotated.permute(0, 2, 3, 1)  # back to (B, Z, Y, X)
 
-    cos_t, sin_t = torch.cos(theta_rad), torch.sin(theta_rad)
-    affine_matrix = torch.stack(
-        [cos_t, -sin_t, torch.zeros_like(theta), sin_t, cos_t, torch.zeros_like(theta)], dim=-1
-    ).view(-1, 2, 3)
 
-    mags = mags.permute(3, 0, 1, 2)
+def tv_loss_vol_sq(obj: torch.Tensor) -> torch.Tensor:
+    """Squared-anisotropic volume TV: sum of squared forward differences.
 
-    def transform_slice(mag_slice):
-        grid = F.affine_grid(affine_matrix, mag_slice.unsqueeze(0).shape, align_corners=False)
-        return F.grid_sample(mag_slice.unsqueeze(0), grid, mode=mode, align_corners=False).squeeze(
-            0
-        )
+    Computes ``Σ (Δd)² + Σ (Δh)² + Σ (Δw)²`` over the three trailing
+    spatial dims, leaving any leading channel/batch axes intact (they are
+    included in the sum). This is the unnormalized ``tv_vol`` regularizer;
+    callers apply their own ``weight / numel`` scaling.
+
+    When the optional ``quantem-cuda`` package is installed
+    (``pip install quantem[cuda]``), the tensor is on a CUDA device, and the
+    ``use_cuda_kernels`` config option is true (default), this dispatches to
+    the fused CUDA forward/backward kernel — identical math, one kernel
+    launch instead of several large intermediates.
+
+    Args:
+        obj: Tensor of shape ``[..., D, H, W]`` (ndim >= 3).
 
-    rotated_mags = torch.vmap(transform_slice)(mags)
-    return rotated_mags.permute(1, 2, 3, 0)
+    Returns:
+        0-dim tensor on the same device as ``obj``; differentiable.
+    """
+    if (
+        obj.is_cuda
+        and obj.dtype == torch.float32
+        and config.get("has_quantem_cuda")
+        and config.get("use_cuda_kernels", default=True)
+    ):
+        from quantem.cuda.core import tv_loss_sq_3d
+
+        return tv_loss_sq_3d(obj)
+
+    tv_d = torch.pow(obj[..., 1:, :, :] - obj[..., :-1, :, :], 2).sum()
+    tv_h = torch.pow(obj[..., :, 1:, :] - obj[..., :, :-1, :], 2).sum()
+    tv_w = torch.pow(obj[..., :, :, 1:] - obj[..., :, :, :-1], 2).sum()
+    return tv_d + tv_h + tv_w
 
 
 def tv_loss_1d(x: torch.Tensor, reduction: str = "mean") -> torch.Tensor:
diff --git a/tests/diffractive_imaging/test_tv_cuda_dispatch.py b/tests/diffractive_imaging/test_tv_cuda_dispatch.py
new file mode 100644
index 000000000..e320ceb31
--- /dev/null
+++ b/tests/diffractive_imaging/test_tv_cuda_dispatch.py
@@ -0,0 +1,103 @@
+"""Tests for ``ObjectConstraints._calc_tv_loss``'s optional quantem-cuda dispatch.
+
+The constraint must produce the same loss and gradients on every path: pure
+torch on CPU, pure torch on GPU (kill-switch or quantem-cuda absent), and
+the fused L1 kernel when quantem-cuda is installed. The dispatch itself is
+asserted by monkeypatching the kernel entry point, so these tests are
+meaningful both with and without quantem-cuda in the environment.
+"""
+
+import pytest
+import torch
+
+from quantem.core import config
+from quantem.diffractive_imaging.object_models import ObjectConstraints
+
+requires_gpu = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a CUDA device")
+requires_quantem_cuda = pytest.mark.skipif(
+    not config.get("has_quantem_cuda"), reason="requires quantem-cuda"
+)
+
+
+class _Host:
+    """Minimal stand-in providing what _calc_tv_loss needs from self."""
+
+    def __init__(self, device):
+        self.device = device
+
+    _get_zero_loss_tensor = ObjectConstraints._get_zero_loss_tensor
+    _calc_tv_loss = ObjectConstraints._calc_tv_loss
+
+
+def _phase(device, shape=(8, 24, 20), seed=0, requires_grad=False):
+    gen = torch.Generator().manual_seed(seed)
+    arr = torch.rand(shape, generator=gen, dtype=torch.float32)
+    return arr.to(device).requires_grad_(requires_grad)
+
+
+WEIGHTS = [(5.0, 0.1), (5.0, 0.0), (0.0, 0.3), (0.0, 0.0)]
+
+
+@requires_gpu
+@pytest.mark.parametrize("weight", WEIGHTS)
+def test_gpu_matches_cpu_reference(weight):
+    expected = _Host("cpu")._calc_tv_loss(_phase("cpu"), weight)
+    actual = _Host("cuda")._calc_tv_loss(_phase("cuda"), weight)
+    torch.testing.assert_close(actual.cpu(), expected, rtol=1e-5, atol=1e-7)
+
+
+@requires_gpu
+@pytest.mark.parametrize("weight", WEIGHTS[:3])
+def test_gpu_grads_match_cpu_reference(weight):
+    a_cpu = _phase("cpu", requires_grad=True)
+    _Host("cpu")._calc_tv_loss(a_cpu, weight).backward()
+    a_gpu = _phase("cuda", requires_grad=True)
+    _Host("cuda")._calc_tv_loss(a_gpu, weight).backward()
+    torch.testing.assert_close(a_gpu.grad.cpu(), a_cpu.grad, rtol=1e-4, atol=1e-7)
+
+
+@requires_gpu
+@pytest.mark.parametrize("shape", [(1, 24, 20), (24, 20)])
+def test_degenerate_and_2d_match_cpu(shape):
+    # num_slices == 1 (weight[0] zeroed upstream, as get_tv_loss does) and
+    # plain 2-D arrays keep exact parity whichever path runs.
+    weight = (0.0, 0.3)
+    expected = _Host("cpu")._calc_tv_loss(_phase("cpu", shape=shape), weight)
+    actual = _Host("cuda")._calc_tv_loss(_phase("cuda", shape=shape), weight)
+    torch.testing.assert_close(actual.cpu(), expected, rtol=1e-5, atol=1e-7)
+
+
+@requires_gpu
+@requires_quantem_cuda
+def test_dispatches_to_kernel(monkeypatch):
+    import quantem.cuda.core
+
+    calls = []
+    real = quantem.cuda.core.tv_loss_l1_3d
+
+    def spy(volume):
+        calls.append(volume.shape)
+        return real(volume)
+
+    monkeypatch.setattr(quantem.cuda.core, "tv_loss_l1_3d", spy)
+    _Host("cuda")._calc_tv_loss(_phase("cuda"), (5.0, 0.1))
+    assert len(calls) == 1
+
+
+@requires_gpu
+@requires_quantem_cuda
+def test_kill_switch_forces_torch_path(monkeypatch):
+    import quantem.cuda.core
+
+    def boom(volume):
+        raise AssertionError("kernel should not be called with use_cuda_kernels=False")
+
+    monkeypatch.setattr(quantem.cuda.core, "tv_loss_l1_3d", boom)
+    expected = _Host("cpu")._calc_tv_loss(_phase("cpu"), (5.0, 0.1))
+    # config.set lacks __exit__, so restore explicitly rather than via `with`.
+    config.set({"use_cuda_kernels": False})
+    try:
+        actual = _Host("cuda")._calc_tv_loss(_phase("cuda"), (5.0, 0.1))
+    finally:
+        config.set({"use_cuda_kernels": True})
+    torch.testing.assert_close(actual.cpu(), expected, rtol=1e-5, atol=1e-7)
diff --git a/tests/ml/test_fused_hidden_mlp.py b/tests/ml/test_fused_hidden_mlp.py
new file mode 100644
index 000000000..b60120cc4
--- /dev/null
+++ b/tests/ml/test_fused_hidden_mlp.py
@@ -0,0 +1,190 @@
+"""Dispatch and model parity for the opt-in cuBLASLt sigma head."""
+
+import copy
+import sys
+from types import ModuleType
+
+import pytest
+import torch
+from torch import nn
+
+from quantem.core import config
+from quantem.core.ml.models import kplanes as kplanes_module
+from quantem.core.ml.models.kplanes import FusedHiddenMLP, KPlanesTILTED
+
+
+def _head():
+    return FusedHiddenMLP(
+        nn.Linear(12, 8),
+        nn.ReLU(inplace=True),
+        nn.Linear(8, 8),
+        nn.ReLU(inplace=True),
+        nn.Linear(8, 1),
+    )
+
+
+@pytest.fixture(autouse=True)
+def reset_fused_mlp_fallback_state():
+    kplanes_module._unsupported_fused_mlp_shapes.clear()
+    kplanes_module._warned_fused_mlp_reasons.clear()
+    yield
+    kplanes_module._unsupported_fused_mlp_shapes.clear()
+    kplanes_module._warned_fused_mlp_reasons.clear()
+
+
+@pytest.mark.parametrize("value", [None, "", "1"])
+def test_fused_mlp_is_default_on_and_only_zero_disables(monkeypatch, value):
+    if value is None:
+        monkeypatch.delenv("QUANTEM_FUSED_MLP", raising=False)
+    else:
+        monkeypatch.setenv("QUANTEM_FUSED_MLP", value)
+    assert kplanes_module._fused_mlp_enabled()
+
+    monkeypatch.setenv("QUANTEM_FUSED_MLP", "0")
+    assert not kplanes_module._fused_mlp_enabled()
+
+
+def test_default_on_still_falls_back_for_unsupported_inputs(monkeypatch):
+    monkeypatch.delenv("QUANTEM_FUSED_MLP", raising=False)
+    head = _head()
+    inputs = torch.randn(7, 12)
+    expected = nn.Sequential.forward(head, inputs)
+    torch.testing.assert_close(head(inputs), expected)
+
+
+def test_unsupported_device_and_autocast_off_fall_back(monkeypatch):
+    monkeypatch.setenv("QUANTEM_FUSED_MLP", "1")
+    head = _head()
+    inputs = torch.randn(7, 12)
+    expected = nn.Sequential.forward(head, inputs)
+    torch.testing.assert_close(head(inputs), expected)
+
+
+def test_missing_extension_capability_falls_back(monkeypatch):
+    monkeypatch.setenv("QUANTEM_FUSED_MLP", "1")
+    head = _head()
+    inputs = torch.randn(7, 12)
+    expected = nn.Sequential.forward(head, inputs)
+    monkeypatch.setitem(sys.modules, "quantem.cuda.core.ml", ModuleType("quantem.cuda.core.ml"))
+    monkeypatch.setattr(FusedHiddenMLP, "_can_fuse", lambda self, tensor: True)
+    with pytest.warns(RuntimeWarning, match="memoizing the eager fallback") as records:
+        torch.testing.assert_close(head(inputs), expected)
+        torch.testing.assert_close(head(inputs), expected)
+    assert len(records) == 1
+
+
+def test_available_extension_is_selected_without_changing_parameter_keys(monkeypatch):
+    monkeypatch.setenv("QUANTEM_FUSED_MLP", "1")
+    head = _head()
+    inputs = torch.randn(7, 12)
+    calls = []
+    cuda_ml = ModuleType("quantem.cuda.core.ml")
+
+    def fake_fused(x, w1, b1, w2, b2, w3, b3):
+        calls.append((x, w1, b1, w2, b2, w3, b3))
+        return torch.full((x.shape[0], w3.shape[0]), 4.0)
+
+    cuda_ml.fused_hidden_mlp = fake_fused
+    monkeypatch.setitem(sys.modules, "quantem.cuda.core.ml", cuda_ml)
+    monkeypatch.setattr(FusedHiddenMLP, "_can_fuse", lambda self, tensor: True)
+
+    output = head(inputs)
+    assert torch.equal(output, torch.full((7, 1), 4.0))
+    assert len(calls) == 1
+    assert set(head.state_dict()) == {
+        "0.weight",
+        "0.bias",
+        "2.weight",
+        "2.bias",
+        "4.weight",
+        "4.bias",
+    }
+
+
+def test_torch_compile_uses_noninterfering_fallback(monkeypatch):
+    monkeypatch.setenv("QUANTEM_FUSED_MLP", "1")
+    head = _head()
+    inputs = torch.randn(7, 12, requires_grad=True)
+    expected = head(inputs)
+    compiled = torch.compile(head, backend="eager")
+    actual = compiled(inputs)
+    torch.testing.assert_close(actual, expected)
+    actual.sum().backward()
+    assert inputs.grad is not None
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+def test_cuda_fp32_autocast_off_falls_back(monkeypatch):
+    monkeypatch.setenv("QUANTEM_FUSED_MLP", "1")
+    head = _head().cuda()
+    inputs = torch.randn(7, 12, device="cuda")
+    expected = nn.Sequential.forward(head, inputs)
+    torch.testing.assert_close(head(inputs), expected)
+
+
+requires_fused_cuda = pytest.mark.skipif(
+    not torch.cuda.is_available() or not config.get("has_quantem_cuda"),
+    reason="requires CUDA and quantem-cuda",
+)
+
+
+@requires_fused_cuda
+def test_kplanes_tilted_forward_backward_parity(monkeypatch):
+    import quantem.cuda.core.ml as cuda_ml
+
+    torch.manual_seed(723)
+    eager_model = KPlanesTILTED(
+        M_features=8,
+        T=1,
+        resolution=(8, 8, 8),
+        multiscale_res_multipliers=(1, 1, 1),
+        density_activation=lambda value: value,
+        use_hybrid_mlp=True,
+        hybrid_hidden_dim=16,
+        hybrid_num_layers=2,
+    ).cuda()
+    fused_model = copy.deepcopy(eager_model)
+    truth_model = copy.deepcopy(eager_model)
+    eager_pts = (torch.rand((37, 3), device="cuda") * 2 - 1).requires_grad_()
+    fused_pts = eager_pts.detach().clone().requires_grad_()
+    truth_pts = eager_pts.detach().clone().requires_grad_()
+    upstream = torch.randn((37, 1), device="cuda", dtype=torch.bfloat16)
+
+    monkeypatch.setenv("QUANTEM_FUSED_MLP", "0")
+    with torch.autocast("cuda", dtype=torch.bfloat16):
+        eager_out = eager_model(eager_pts)
+    eager_out.backward(upstream)
+
+    truth_out = truth_model(truth_pts)
+    truth_out.backward(upstream.float())
+
+    real_fused = cuda_ml.fused_hidden_mlp
+    calls = []
+
+    def spy(*args):
+        calls.append(args[0].shape)
+        return real_fused(*args)
+
+    monkeypatch.setattr(cuda_ml, "fused_hidden_mlp", spy)
+    monkeypatch.setenv("QUANTEM_FUSED_MLP", "1")
+    with torch.autocast("cuda", dtype=torch.bfloat16):
+        fused_out = fused_model(fused_pts)
+    fused_out.backward(upstream)
+
+    assert calls == [torch.Size((37, 24))]
+    torch.testing.assert_close(fused_out, eager_out, rtol=8e-3, atol=1e-3)
+    eager_grads = (eager_pts.grad, *(parameter.grad for parameter in eager_model.parameters()))
+    fused_grads = (fused_pts.grad, *(parameter.grad for parameter in fused_model.parameters()))
+    truth_grads = (truth_pts.grad, *(parameter.grad for parameter in truth_model.parameters()))
+    # At M=37 the measured fused/eager mean-error ratios against fp32 truth
+    # were 1.11 for dx and 1.34 for dW1; allow reduction-order variation while
+    # requiring the fused path to remain close to eager's fp32 error envelope.
+    for fused_grad, eager_grad, truth_grad in zip(fused_grads, eager_grads, truth_grads):
+        eager_error = (eager_grad.float() - truth_grad).abs()
+        fused_error = (fused_grad.float() - truth_grad).abs()
+        assert torch.mean(fused_error) <= torch.mean(eager_error) * 1.5 + 1e-7
+        truth_scale = truth_grad.abs().max()
+        bf16_eps_at_truth_scale = truth_scale * 2**-8
+        assert torch.max(fused_error) <= (
+            torch.max(eager_error) * 1.5 + 2 * bf16_eps_at_truth_scale
+        )
diff --git a/tests/ml/test_inr.py b/tests/ml/test_inr.py
new file mode 100644
index 000000000..c78bc1af3
--- /dev/null
+++ b/tests/ml/test_inr.py
@@ -0,0 +1,40 @@
+"""Tests for ``quantem.core.ml.inr`` Siren / HSiren construction."""
+
+import torch
+
+from quantem.core.ml.inr import HSiren, Siren
+
+
+def _first_layer_weight(winner_seed) -> torch.Tensor:
+    # Fix the global RNG so the base SineLayer init is identical across builds;
+    # only the winner-initialization perturbation varies with winner_seed.
+    torch.manual_seed(0)
+    model = HSiren(hidden_layers=1, hidden_features=8, winner_initialization=winner_seed)
+    return model.net[0].linear.weight.detach().clone()
+
+
+class TestWinnerInitialization:
+    def test_same_seed_is_reproducible(self):
+        assert torch.equal(_first_layer_weight(72), _first_layer_weight(72))
+
+    def test_different_seeds_differ(self):
+        """Regression: the seeded torch.Generator was created but never used
+        (torch.randn_like ignores generators), so every winner seed produced
+        the same perturbation."""
+        assert not torch.equal(_first_layer_weight(72), _first_layer_weight(73))
+
+    def test_true_uses_default_seed(self):
+        assert torch.equal(_first_layer_weight(True), _first_layer_weight(42))
+
+    def test_disabled_adds_no_perturbation(self):
+        torch.manual_seed(0)
+        base = Siren(hidden_layers=1, hidden_features=8, winner_initialization=False)
+        torch.manual_seed(0)
+        again = Siren(hidden_layers=1, hidden_features=8, winner_initialization=False)
+        assert torch.equal(base.net[0].linear.weight.detach(), again.net[0].linear.weight.detach())
+
+
+def test_forward_shape():
+    model = HSiren(hidden_layers=1, hidden_features=8)
+    out = model(torch.rand(5, 3))
+    assert out.shape == (5, 1)
diff --git a/tests/ml/test_kplanes.py b/tests/ml/test_kplanes.py
new file mode 100644
index 000000000..78cf6f98f
--- /dev/null
+++ b/tests/ml/test_kplanes.py
@@ -0,0 +1,103 @@
+"""Tests for ``quantem.core.ml.models.kplanes`` construction guards."""
+
+import io
+
+import pytest
+import torch
+
+from quantem.core.ml.models.kplanes import KPlanes, KPlanesTILTED
+from quantem.core.ml.models.so3params import SO3ParamR9SVD
+
+
+def test_r9svd_as_matrix_stays_fp32_under_autocast():
+    so3 = SO3ParamR9SVD(T=2)
+
+    with torch.autocast(device_type="cpu", dtype=torch.bfloat16, enabled=True):
+        rotation_matrices = so3.as_matrix()
+
+    assert rotation_matrices.dtype == torch.float32
+
+
+class TestResolutionValidation:
+    def test_isotropic_builds(self):
+        model = KPlanes(M_features=2, resolution=(16, 16, 16))
+        assert model.grids[0].shape == (3, 2, 16, 16)
+
+    def test_anisotropic_raises(self):
+        """Regression: the plane grids are allocated as (res[1], res[0]) for all three
+        axis pairs, ignoring res[2] -- anisotropic resolutions silently gave the
+        XZ/YZ planes the wrong grid along z instead of erroring."""
+        with pytest.raises(ValueError, match="isotropic"):
+            KPlanes(M_features=2, resolution=(16, 16, 8))
+        with pytest.raises(ValueError, match="isotropic"):
+            KPlanesTILTED(M_features=2, T=2, resolution=(16, 8, 16))
+
+
+@pytest.mark.parametrize(
+    ("model_type", "kwargs"),
+    [
+        (KPlanes, {}),
+        (KPlanesTILTED, {"T": 2}),
+    ],
+)
+def test_grid_parameters_are_channels_last(model_type, kwargs):
+    model = model_type(M_features=2, resolution=(8, 8, 8), **kwargs)
+
+    for plane in model.grids:
+        assert plane.is_contiguous(memory_format=torch.channels_last)
+        assert plane.permute(0, 2, 3, 1).is_contiguous()
+
+
+def test_whole_module_load_reformats_legacy_grids():
+    model = KPlanesTILTED(
+        M_features=2,
+        T=2,
+        resolution=(8, 8, 8),
+        density_activation=torch.relu,
+    )
+    for plane in model.grids:
+        plane.data = plane.data.contiguous()
+
+    checkpoint = io.BytesIO()
+    torch.save(model, checkpoint)
+    checkpoint.seek(0)
+    restored = torch.load(checkpoint, weights_only=False)
+
+    assert all(plane.is_contiguous(memory_format=torch.channels_last) for plane in restored.grids)
+
+
+class TestDefaultHeadConstruction:
+    """Regression: KPlanes(use_hybrid_mlp=False) -- the constructor default --
+    built no sigma_net at all, so forward / get_params / ObjectTensorDecomp
+    .from_model crashed with AttributeError. KPlanesTILTED and CPTilted both
+    fall back to a linear head; plain KPlanes must do the same."""
+
+    def test_default_get_params(self):
+        model = KPlanes(M_features=2, resolution=(8, 8, 8))
+        params = model.get_params()
+        assert set(params) == set(model.param_keys)
+        assert all(len(v) > 0 for v in params.values())
+
+    def test_default_forward(self):
+        import torch
+
+        model = KPlanes(M_features=2, resolution=(8, 8, 8))
+        out = model(torch.rand(5, 3) * 2 - 1)
+        assert out.shape == (5, 1)
+        assert torch.isfinite(out).all()
+
+
+def test_tilted_forward_is_dynamo_traceable():
+    """The public ``grids`` child must not be hidden behind a broken property.
+
+    ``nn.Module.__setattr__`` registers the ParameterList as ``grids``.  Eager
+    lookup tolerates a same-named property whose getter reads missing ``_grids``,
+    but Dynamo traces that getter and turns its AttributeError into Unsupported.
+    """
+    model = KPlanesTILTED(M_features=2, T=2, resolution=(4, 4, 4))
+    coords = torch.rand(8, 3) * 2 - 1
+
+    eager = model(coords)
+    compiled = torch.compile(model, backend="eager", fullgraph=True)(coords)
+
+    torch.testing.assert_close(compiled, eager)
diff --git a/tests/ml/test_kplanes_cuda_dispatch.py b/tests/ml/test_kplanes_cuda_dispatch.py
new file mode 100644
index 000000000..e34567b00
--- /dev/null
+++ b/tests/ml/test_kplanes_cuda_dispatch.py
@@ -0,0 +1,131 @@
+"""Tests for ``interpolate_ms_features_tilted``'s optional quantem-cuda dispatch.
+
+The function must produce the same features and gradients on every path:
+pure torch on CPU, pure torch on GPU (kill-switch or quantem-cuda absent),
+and the fused CUDA kernel when quantem-cuda is installed. The dispatch
+itself is asserted by monkeypatching the kernel entry point, so these tests
+are meaningful both with and without quantem-cuda in the environment.
+"""
+
+import pytest
+import torch
+from torch import nn
+
+from quantem.core import config
+from quantem.core.ml.models.kplanes import KPlanesTILTED, interpolate_ms_features_tilted
+
+requires_gpu = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a CUDA device")
+requires_quantem_cuda = pytest.mark.skipif(
+    not config.get("has_quantem_cuda"), reason="requires quantem-cuda"
+)
+
+
+def _inputs(device, requires_grad=False, seed=0, B=64, T=3, C=5, scales=(9, 17)):
+    gen = torch.Generator().manual_seed(seed)
+    pts = torch.rand(B, 3, generator=gen) * 2.2 - 1.1  # some points outside [-1, 1]
+    rotations = torch.rand(T, 3, 3, generator=gen) * 2 - 1
+    grids = nn.ParameterList(
+        nn.Parameter(torch.rand(3 * T, C, s, s, generator=gen) * 0.4 + 0.1) for s in scales
+    )
+    pts = pts.to(device).requires_grad_(requires_grad)
+    rotations = rotations.to(device).requires_grad_(requires_grad)
+    grids = grids.to(device)
+    if not requires_grad:
+        for g in grids:
+            g.requires_grad_(False)
+    return pts, rotations, grids
+
+
+@requires_gpu
+def test_gpu_matches_cpu_reference():
+    pts_c, rot_c, grids_c = _inputs("cpu")
+    expected = interpolate_ms_features_tilted(pts_c, grids_c, rot_c)
+    pts_g, rot_g, grids_g = _inputs("cuda")
+    actual = interpolate_ms_features_tilted(pts_g, grids_g, rot_g).cpu()
+    torch.testing.assert_close(actual, expected, rtol=1e-4, atol=1e-5)
+
+
+@requires_gpu
+def test_gpu_grads_match_cpu_reference():
+    pts_c, rot_c, grids_c = _inputs("cpu", requires_grad=True)
+    interpolate_ms_features_tilted(pts_c, grids_c, rot_c).square().sum().backward()
+    pts_g, rot_g, grids_g = _inputs("cuda", requires_grad=True)
+    interpolate_ms_features_tilted(pts_g, grids_g, rot_g).square().sum().backward()
+    torch.testing.assert_close(pts_g.grad.cpu(), pts_c.grad, rtol=1e-3, atol=1e-5)
+    torch.testing.assert_close(rot_g.grad.cpu(), rot_c.grad, rtol=1e-3, atol=1e-5)
+    for gg, gc in zip(grids_g, grids_c):
+        torch.testing.assert_close(gg.grad.cpu(), gc.grad, rtol=1e-3, atol=1e-5)
+
+
+@requires_gpu
+@requires_quantem_cuda
+def test_dispatches_to_kernel_per_scale(monkeypatch):
+    import quantem.cuda.core.ml
+
+    calls = []
+    real = quantem.cuda.core.ml.kplanes_tilted_fuse
+
+    def spy(pts, rotations, plane):
+        calls.append(plane.shape)
+        return real(pts, rotations, plane)
+
+    monkeypatch.setattr(quantem.cuda.core.ml, "kplanes_tilted_fuse", spy)
+    pts, rot, grids = _inputs("cuda")
+    interpolate_ms_features_tilted(pts, grids, rot)
+    assert len(calls) == len(grids)
+
+
+@requires_gpu
+@requires_quantem_cuda
+def test_kill_switch_forces_torch_path(monkeypatch):
+    import quantem.cuda.core.ml
+
+    def boom(pts, rotations, plane):
+        raise AssertionError("kernel should not be called with use_cuda_kernels=False")
+
+    monkeypatch.setattr(quantem.cuda.core.ml, "kplanes_tilted_fuse", boom)
+    pts, rot, grids = _inputs("cuda")
+    pts_c, rot_c, grids_c = _inputs("cpu")
+    expected = interpolate_ms_features_tilted(pts_c, grids_c, rot_c)
+    # config.set lacks __exit__, so restore explicitly rather than via `with`.
+    config.set({"use_cuda_kernels": False})
+    try:
+        actual = interpolate_ms_features_tilted(pts, grids, rot).cpu()
+    finally:
+        config.set({"use_cuda_kernels": True})
+    torch.testing.assert_close(actual, expected, rtol=1e-4, atol=1e-5)
+
+
+@requires_gpu
+def test_non_fp32_takes_torch_path():
+    pts, rot, grids = _inputs("cuda")
+    pts64 = pts.double()
+    rot64 = rot.double()
+    grids64 = nn.ParameterList(nn.Parameter(g.double()) for g in grids)
+    out = interpolate_ms_features_tilted(pts64, grids64, rot64)
+    assert out.dtype == torch.float64
+
+
+@requires_gpu
+@requires_quantem_cuda
+def test_three_scale_features_enter_sigma_net_in_autocast_dtype(monkeypatch):
+    monkeypatch.setenv("QUANTEM_KPLANES_MS_BF16_OUT", "1")
+    model = KPlanesTILTED(
+        M_features=2,
+        T=2,
+        resolution=(8, 8, 8),
+        multiscale_res_multipliers=(0.5, 0.75, 1.0),
+    ).cuda()
+    coords = torch.rand((16, 3), device="cuda", dtype=torch.float32) * 2 - 1
+    input_dtypes = []
+    handle = model.sigma_net.register_forward_pre_hook(
+        lambda _module, inputs: input_dtypes.append(inputs[0].dtype)
+    )
+    try:
+        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
+            model(coords)
+        model(coords)
+    finally:
+        handle.remove()
+
+    assert input_dtypes == [torch.bfloat16, torch.float32]
diff --git a/tests/ml/test_kplanes_cuda_ms_dispatch.py b/tests/ml/test_kplanes_cuda_ms_dispatch.py
new file mode 100644
index 000000000..18d412794
--- /dev/null
+++ b/tests/ml/test_kplanes_cuda_ms_dispatch.py
@@ -0,0 +1,149 @@
+"""CPU-only dispatch coverage for the optional multiscale K-Planes CUDA op."""
+
+import sys
+from types import ModuleType, SimpleNamespace
+
+import pytest
+import torch
+
+from quantem.core import config
+from quantem.core.ml.models.kplanes import interpolate_ms_features_tilted
+
+
+def _inputs(num_grids=3):
+    pts = SimpleNamespace(is_cuda=True, dtype=torch.float32, shape=(32, 3))
+    rotations = SimpleNamespace(dtype=torch.float32, shape=(2, 3, 3))
+    grids = [torch.empty((6, 4, 1, 1)) for _ in range(num_grids)]
+    return pts, rotations, grids
+
+
+@pytest.fixture
+def mocked_cuda_ml(monkeypatch):
+    calls = {"single": [], "ms": [], "ms_tv": []}
+    cuda_ml = ModuleType("quantem.cuda.core.ml")
+
+    def single(pts, rotations, grid):
+        calls["single"].append((pts, rotations, grid))
+        return torch.zeros((pts.shape[0], rotations.shape[0] * grid.shape[1]))
+
+    def multiscale(*args):
+        calls["ms"].append((args[2:5], args[5:8]))
+        width = args[1].shape[0] * sum(grid.shape[1] for grid in args[2:5])
+        return torch.zeros((args[0].shape[0], width))
+
+    def multiscale_tv(*args):
+        calls["ms_tv"].append((args[2:5], args[5:8]))
+        width = args[1].shape[0] * sum(grid.shape[1] for grid in args[2:5])
+        return torch.zeros((args[0].shape[0], width)), torch.tensor(2.5)
+
+    cuda_ml.kplanes_tilted_fuse = single
+    cuda_ml._kplanes_tilted_fuse_builtin = single
+    cuda_ml.kplanes_tilted_fuse_ms = multiscale
+    cuda_ml.kplanes_tilted_fuse_ms_tv = multiscale_tv
+
+    cuda = ModuleType("quantem.cuda")
+    cuda.__path__ = []
+    cuda_core = ModuleType("quantem.cuda.core")
+    cuda_core.__path__ = []
+    cuda.core = cuda_core
+    cuda_core.ml = cuda_ml
+    monkeypatch.setattr(sys.modules["quantem"], "cuda", cuda, raising=False)
+    monkeypatch.setitem(sys.modules, "quantem.cuda", cuda)
+    monkeypatch.setitem(sys.modules, "quantem.cuda.core", cuda_core)
+    monkeypatch.setitem(sys.modules, "quantem.cuda.core.ml", cuda_ml)
+    monkeypatch.setattr(config, "get", lambda key, default=None: True)
+    return cuda_ml, calls
+
+
+def test_three_scales_dispatch_to_one_multiscale_op(mocked_cuda_ml):
+    _, calls = mocked_cuda_ml
+    pts, rotations, grids = _inputs()
+    output = interpolate_ms_features_tilted(pts, grids, rotations, scale_gates=(0.25, 0.75, 1.0))
+
+    assert output.shape == (pts.shape[0], rotations.shape[0] * 4 * 3)
+    assert len(calls["ms"]) == 1
+    called_grids, called_scales = calls["ms"][0]
+    assert all(actual is expected for actual, expected in zip(called_grids, grids))
+    assert called_scales == (0.25, 0.75, 1.0)
+    assert calls["single"] == []
+
+
+def test_none_scale_gates_dispatch_as_unit_gates(mocked_cuda_ml):
+    _, calls = mocked_cuda_ml
+    pts, rotations, grids = _inputs()
+
+    interpolate_ms_features_tilted(pts, grids, rotations, scale_gates=None)
+
+    assert len(calls["ms"]) == 1
+    assert calls["ms"][0][1] == (1.0, 1.0, 1.0)
+    assert calls["single"] == []
+
+
+def test_plane_tv_defaults_to_combined_capability(mocked_cuda_ml, monkeypatch):
+    _, calls = mocked_cuda_ml
+    monkeypatch.delenv("QUANTEM_KPLANES_MS_TV_FUSED", raising=False)
+    pts, rotations, grids = _inputs()
+
+    features, tv = interpolate_ms_features_tilted(
+        pts, grids, rotations, scale_gates=(0.25, 0.75, 1.0), include_plane_tv=True
+    )
+
+    assert features.shape == (pts.shape[0], rotations.shape[0] * 4 * 3)
+    assert tv.item() == 2.5
+    assert len(calls["ms_tv"]) == 1
+    assert calls["ms"] == []
+
+
+def test_plane_tv_opt_out_uses_multiscale_without_aux(mocked_cuda_ml, monkeypatch):
+    _, calls = mocked_cuda_ml
+    monkeypatch.setenv("QUANTEM_KPLANES_MS_TV_FUSED", "0")
+    pts, rotations, grids = _inputs()
+
+    interpolate_ms_features_tilted(pts, grids, rotations, include_plane_tv=True)
+
+    assert len(calls["ms"]) == 1
+    assert calls["ms_tv"] == []
+
+
+@pytest.mark.parametrize("num_grids", [2, 4])
+def test_non_three_scale_counts_use_per_level_fallback(mocked_cuda_ml, num_grids):
+    _, calls = mocked_cuda_ml
+    pts, rotations, grids = _inputs(num_grids)
+
+    interpolate_ms_features_tilted(pts, grids, rotations)
+
+    assert len(calls["single"]) == len(grids)
+    assert all(call[2] is grid for call, grid in zip(calls["single"], grids))
+    assert calls["ms"] == []
+
+
+def test_missing_multiscale_op_uses_per_level_fallback(mocked_cuda_ml, monkeypatch):
+    cuda_ml, calls = mocked_cuda_ml
+    monkeypatch.setattr(cuda_ml, "kplanes_tilted_fuse_ms", None)
+    pts, rotations, grids = _inputs()
+
+    interpolate_ms_features_tilted(pts, grids, rotations)
+
+    assert len(calls["single"]) == len(grids)
+    assert all(call[2] is grid for call, grid in zip(calls["single"], grids))
+    assert calls["ms"] == []
+
+
+def test_overridden_single_level_op_uses_per_level_fallback(mocked_cuda_ml, monkeypatch):
+    cuda_ml, calls = mocked_cuda_ml
+    instrumented_calls = []
+    builtin = cuda_ml.kplanes_tilted_fuse
+
+    def instrumented(*args):
+        instrumented_calls.append(args)
+        return builtin(*args)
+
+    monkeypatch.setattr(cuda_ml, "kplanes_tilted_fuse", instrumented)
+    pts, rotations, grids = _inputs()
+
+    interpolate_ms_features_tilted(pts, grids, rotations)
+
+    assert len(instrumented_calls) == len(grids)
+    assert all(call[2] is grid for call, grid in zip(instrumented_calls, grids))
+    assert len(calls["single"]) == len(grids)
+    assert calls["ms"] == []
diff --git a/tests/ml/test_kplanes_ms_tv_optimizer.py b/tests/ml/test_kplanes_ms_tv_optimizer.py
new file mode 100644
index 000000000..dc0109702
--- /dev/null
+++ b/tests/ml/test_kplanes_ms_tv_optimizer.py
@@ -0,0 +1,72 @@
+"""End-to-end optimizer parity for combined multiscale interpolation and plane TV."""
+
+import copy
+
+import pytest
+import torch
+
+from quantem.core import config
+from quantem.core.ml.models.kplanes import KPlanesTILTED
+
+requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+
+
+@requires_cuda
+def test_combined_plane_tv_matches_separate_path_over_optimizer_steps(monkeypatch):
+    try:
+        import quantem.cuda.core.ml as cuda_ml
+    except (ImportError, RuntimeError):
+        pytest.skip("quantem-cuda is unavailable")
+    if not hasattr(cuda_ml, "kplanes_tilted_fuse_ms_tv"):
+        pytest.skip("combined multiscale/plane-TV capability is unavailable")
+
+    original_get = config.get
+
+    def enabled_cuda_config(key, default=None):
+        if key in {"has_quantem_cuda", "use_cuda_kernels"}:
+            return True
+        return original_get(key, default=default)
+
+    monkeypatch.setattr(config, "get", enabled_cuda_config)
+    monkeypatch.setenv("QUANTEM_KPLANES_MS_TV_FUSED", "1")
+    torch.manual_seed(41)
+    separate = KPlanesTILTED(
+        T=1,
+        M_features=2,
+        resolution=(5, 5, 5),
+        multiscale_res_multipliers=(1, 2, 3),
+    ).cuda()
+    combined = copy.deepcopy(separate)
+    coords = torch.tensor([[0.125, -0.25, 0.375]], device="cuda")
+    tv_weight = 0.0375
+    separate_optim = torch.optim.SGD(separate.parameters(), lr=2e-4)
+    combined_optim = torch.optim.SGD(combined.parameters(), lr=2e-4)
+
+    for _ in range(3):
+        separate_optim.zero_grad(set_to_none=True)
+        combined_optim.zero_grad(set_to_none=True)
+        separate._plane_tv_fusion_requested = False
+        combined._plane_tv_fusion_requested = True
+
+        separate_density = separate(coords)
+        combined_density, combined_tv = combined(coords)
+        separate_tv = cuda_ml.plane_tv_loss(*separate.grids)
+        separate_loss = separate_density.square().mean() + tv_weight * separate_tv
+        combined_loss = combined_density.square().mean() + tv_weight * combined_tv
+        torch.testing.assert_close(combined_loss, separate_loss, rtol=2e-6, atol=2e-7)
+
+        separate_loss.backward()
+        combined_loss.backward()
+        for separate_grid, combined_grid in zip(separate.grids, combined.grids):
+            torch.testing.assert_close(
+                combined_grid.grad, separate_grid.grad, rtol=1e-6, atol=1e-7
+            )
+        for separate_param, combined_param in zip(separate.parameters(), combined.parameters()):
+            torch.testing.assert_close(
+                combined_param.grad, separate_param.grad, rtol=1e-6, atol=1e-7
+            )
+
+        separate_optim.step()
+        combined_optim.step()
+        for separate_param, combined_param in zip(separate.parameters(), combined.parameters()):
+            torch.testing.assert_close(combined_param, separate_param, rtol=1e-6, atol=1e-7)
diff --git a/tests/ml/test_optimizermixin.py b/tests/ml/test_optimizermixin.py
index d6347ce66..8cbc1bc9a 100644
--- a/tests/ml/test_optimizermixin.py
+++ b/tests/ml/test_optimizermixin.py
@@ -356,9 +356,7 @@ def test_parse_default_name_is_none(self):
 from quantem.core.ml.optimizer_mixin import OptimizerMixin  # noqa: E402
 
 torch = pytest.importorskip("torch") if config.get("has_torch") else None
-requires_torch = pytest.mark.skipif(
-    not config.get("has_torch"), reason="requires torch"
-)
+requires_torch = pytest.mark.skipif(not config.get("has_torch"), reason="requires torch")
 
 
 def _param(value=1.0):
@@ -413,7 +411,10 @@ def test_none_optimizer_removes_without_touching_params(self):
     def test_multi_group_pplr_applies_per_group_lr(self):
         model = _FakeModel({"descan": [_param()], "scan_positions": [_param(2.0)]})
         model.set_optimizer(
-            {"descan": OptimizerParams.SGD(lr=1e-2), "scan_positions": OptimizerParams.SGD(lr=1e-3)}
+            {
+                "descan": OptimizerParams.SGD(lr=1e-2),
+                "scan_positions": OptimizerParams.SGD(lr=1e-3),
+            }
         )
         groups = model.optimizer.param_groups
         assert len(groups) == 2
@@ -445,9 +446,118 @@ def test_reset_optimizer_on_unconfigured_model_is_noop(self):
     def test_set_scheduler_base_lr_uses_max_group_lr(self):
         model = _FakeModel({"descan": [_param()], "scan_positions": [_param(2.0)]})
         model.set_optimizer(
-            {"descan": OptimizerParams.SGD(lr=1e-2), "scan_positions": OptimizerParams.SGD(lr=1e-3)}
+            {
+                "descan": OptimizerParams.SGD(lr=1e-2),
+                "scan_positions": OptimizerParams.SGD(lr=1e-3),
+            }
         )
         model.set_scheduler(SchedulerParams.Plateau(), num_iter=10)
         assert model.scheduler is not None
         # Plateau min_lr defaults to base_LR / 20, with base_LR = max group lr (1e-2)
         assert model.scheduler.min_lrs[0] == pytest.approx(1e-2 / 20)
+
+
+class TestSchedulerParamsArePure:
+    """Regression: params() mutated the dataclass, so a shared instance baked in the
+    first call's derived values (min_lr, gamma, bounds, horizons) for every later
+    optimizer or reconstruct() call."""
+
+    def test_plateau_min_lr_follows_base_lr(self):
+        p = SchedulerParams.Plateau()
+        assert p.params(base_LR=1.0)["min_lr"] == pytest.approx(1 / 20)
+        assert p.params(base_LR=10.0)["min_lr"] == pytest.approx(10 / 20)
+        assert p.min_lr is None  # instance untouched
+
+    def test_exponential_gamma_follows_num_iter(self):
+        p = SchedulerParams.Exponential(factor=0.5)
+        g10 = p.params(base_LR=1.0, num_iter=10)["gamma"]
+        g100 = p.params(base_LR=1.0, num_iter=100)["gamma"]
+        assert g10 == pytest.approx(0.5 ** (1 / 10))
+        assert g100 == pytest.approx(0.5 ** (1 / 100))
+        assert p.num_iter is None and p.gamma == 0.9
+
+    def test_cyclic_bounds_follow_base_lr(self):
+        p = SchedulerParams.Cyclic()
+        first = p.params(base_LR=1.0)
+        second = p.params(base_LR=2.0)
+        assert second["base_lr"] == pytest.approx(2 * first["base_lr"])
+        assert second["max_lr"] == pytest.approx(2 * first["max_lr"])
+        assert p.base_lr is None and p.max_lr is None
+
+    def test_linear_and_cosine_follow_num_iter(self):
+        lin = SchedulerParams.Linear()
+        assert lin.params(base_LR=1.0, num_iter=5)["total_iters"] == 5
+        assert lin.params(base_LR=1.0, num_iter=50)["total_iters"] == 50
+        assert lin.total_iters is None
+        cos = SchedulerParams.CosineAnnealing()
+        assert cos.params(base_LR=1.0, num_iter=5)["T_max"] == 5
+        assert cos.params(base_LR=1.0, num_iter=50)["T_max"] == 50
+        assert cos.T_max is None
+
+    def test_explicit_values_still_win(self):
+        p = SchedulerParams.Plateau(min_lr=1e-6)
+        assert p.params(base_LR=1.0)["min_lr"] == 1e-6
+        e = SchedulerParams.Exponential(gamma=0.8)
+        assert e.params(base_LR=1.0, num_iter=10)["gamma"] == 0.8
+
+
+@requires_torch
+class TestReconnectHyperparamAlignment:
+    """Regression: reconnect_optimizer_to_parameters restored per-group
+    hyperparameters by zip-index against get_optimization_parameters(),
+    silently attaching the wrong lr to every group whenever the group dict's
+    order (or membership) changed between optimizer creation and reconnect."""
+
+    def _model(self):
+        model = _FakeModel({"a": [_param()], "b": [_param(2.0)]})
+        model.set_optimizer({"a": OptimizerParams.SGD(lr=1e-2), "b": OptimizerParams.SGD(lr=1e-3)})
+        return model
+
+    def _lrs_by_name(self, model):
+        return {pg["name"]: pg["lr"] for pg in model.optimizer.param_groups}
+
+    def test_groups_carry_names(self):
+        model = self._model()
+        assert self._lrs_by_name(model) == {"a": 1e-2, "b": 1e-3}
+
+    def test_reordered_groups_keep_their_lr(self):
+        model = self._model()
+        # Simulate get_optimization_parameters returning a different order at
+        # reconnect time (group membership/order is not contractual).
+        model._groups = {"b": model._groups["b"], "a": model._groups["a"]}
+        model.reconnect_optimizer_to_parameters()
+        assert self._lrs_by_name(model) == {"a": 1e-2, "b": 1e-3}
+
+    def test_added_group_gets_defaults_others_keep_lr(self):
+        model = self._model()
+        model._groups = dict(model._groups)
+        model._groups["c"] = [_param(3.0)]
+        model.reconnect_optimizer_to_parameters()
+        lrs = self._lrs_by_name(model)
+        assert lrs["a"] == 1e-2 and lrs["b"] == 1e-3
+        assert "c" in lrs  # present, with optimizer defaults
+
+    def test_state_survives_reconnect(self):
+        # Adam: SGD without momentum keeps no per-param state to preserve.
+        model = _FakeModel({"a": [_param()], "b": [_param(2.0)]})
+        model.set_optimizer(
+            {"a": OptimizerParams.Adam(lr=1e-2), "b": OptimizerParams.Adam(lr=1e-3)}
+        )
+        params = [pg["params"][0] for pg in model.optimizer.param_groups]
+        for p in params:
+            p.grad = torch.ones_like(p)
+        model.optimizer.step()
+        assert len(model.optimizer.state) > 0
+        before = {id(p): dict(model.optimizer.state[p]) for p in params}
+        model.reconnect_optimizer_to_parameters()
+        for p in params:
+            assert id(p) in before and p in model.optimizer.state
+
+    def test_legacy_groups_without_names_fall_back_to_index(self):
+        model = self._model()
+        # Optimizers restored from older checkpoints carry no group names.
+        for pg in model.optimizer.param_groups:
+            del pg["name"]
+        model.reconnect_optimizer_to_parameters()
+        lrs = [pg["lr"] for pg in model.optimizer.param_groups]
+        assert lrs == [1e-2, 1e-3]
diff --git a/tests/ml/test_trunc_exp_activation.py b/tests/ml/test_trunc_exp_activation.py
new file mode 100644
index 000000000..e51d865db
--- /dev/null
+++ b/tests/ml/test_trunc_exp_activation.py
@@ -0,0 +1,23 @@
+"""Reference and capability coverage for tomography's trunc-exp activation."""
+
+import torch
+
+from quantem.core.ml.activation_functions import TruncExpActivation, trunc_exp
+
+
+def test_trunc_exp_cpu_forward_and_clamped_backward():
+    values = torch.tensor([-2.0, 0.5, 20.0], requires_grad=True)
+    upstream = torch.tensor([0.25, -0.5, 0.75])
+    output = trunc_exp(values, offset=1.25)
+    output.backward(upstream)
+
+    shifted = values.detach() - 1.25
+    torch.testing.assert_close(output, torch.exp(shifted))
+    torch.testing.assert_close(values.grad, upstream * torch.exp(shifted.clamp(max=15)))
+
+
+def test_trunc_exp_capability_is_explicit_on_function_and_module():
+    activation = TruncExpActivation(offset=1.5)
+    assert trunc_exp.quantem_guarantees_nonnegative is True
+    assert activation.quantem_guarantees_nonnegative is True
+    assert activation.offset == 1.5
diff --git a/tests/tomography/test_cuda_dispatch.py b/tests/tomography/test_cuda_dispatch.py
new file mode 100644
index 000000000..db6e96914
--- /dev/null
+++ b/tests/tomography/test_cuda_dispatch.py
@@ -0,0 +1,91 @@
+"""Tests for ``tv_loss_vol_sq`` and its optional quantem-cuda dispatch.
+
+The helper must produce the same value and gradients on every path: pure
+torch on CPU, pure torch on GPU (kill-switch or quantem-cuda absent), and
+the fused CUDA kernel when quantem-cuda is installed. The dispatch itself is
+asserted by monkeypatching the kernel entry point, so these tests are
+meaningful both with and without quantem-cuda in the environment.
+"""
+
+import pytest
+import torch
+
+from quantem.core import config
+from quantem.tomography.utils import tv_loss_vol_sq
+
+from .conftest import requires_gpu
+
+requires_quantem_cuda = pytest.mark.skipif(
+    not config.get("has_quantem_cuda"), reason="requires quantem-cuda"
+)
+
+
+def tv_vol_sq_ref(obj: torch.Tensor) -> torch.Tensor:
+    tv_d = torch.pow(obj[..., 1:, :, :] - obj[..., :-1, :, :], 2).sum()
+    tv_h = torch.pow(obj[..., :, 1:, :] - obj[..., :, :-1, :], 2).sum()
+    tv_w = torch.pow(obj[..., :, :, 1:] - obj[..., :, :, :-1], 2).sum()
+    return tv_d + tv_h + tv_w
+
+
+@pytest.mark.parametrize("shape", [(7, 6, 5), (2, 7, 6, 5)])
+def test_cpu_matches_reference(shape):
+    obj = torch.rand(shape, generator=torch.Generator().manual_seed(0))
+    torch.testing.assert_close(tv_loss_vol_sq(obj), tv_vol_sq_ref(obj))
+
+
+def test_constant_volume_is_zero():
+    assert tv_loss_vol_sq(torch.ones(8, 8, 8)).item() == 0.0
+
+
+@requires_gpu
+@pytest.mark.parametrize("shape", [(7, 6, 5), (2, 7, 6, 5)])
+def test_gpu_matches_cpu_reference(shape):
+    obj = torch.rand(shape, generator=torch.Generator().manual_seed(1))
+    expected = tv_vol_sq_ref(obj)
+    actual = tv_loss_vol_sq(obj.cuda()).cpu()
+    torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-6)
+
+
+@requires_gpu
+def test_gpu_grad_matches_cpu_reference():
+    obj = torch.rand(6, 7, 8, generator=torch.Generator().manual_seed(2))
+    v_cpu = obj.clone().requires_grad_(True)
+    tv_vol_sq_ref(v_cpu).backward()
+    v_gpu = obj.cuda().requires_grad_(True)
+    tv_loss_vol_sq(v_gpu).backward()
+    torch.testing.assert_close(v_gpu.grad.cpu(), v_cpu.grad, rtol=1e-4, atol=1e-6)
+
+
+@requires_gpu
+@requires_quantem_cuda
+def test_dispatches_to_kernel(monkeypatch):
+    import quantem.cuda.core
+
+    calls = []
+    real = quantem.cuda.core.tv_loss_sq_3d
+
+    def spy(volume):
+        calls.append(volume.shape)
+        return real(volume)
+
+    monkeypatch.setattr(quantem.cuda.core, "tv_loss_sq_3d", spy)
+    tv_loss_vol_sq(torch.rand(4, 4, 4, device="cuda"))
+    assert len(calls) == 1
+
+
+@requires_gpu
+@requires_quantem_cuda
+def test_kill_switch_forces_torch_path(monkeypatch):
+    import quantem.cuda.core
+
+    def boom(volume):
+        raise AssertionError("kernel should not be called with use_cuda_kernels=False")
+
+    monkeypatch.setattr(quantem.cuda.core, "tv_loss_sq_3d", boom)
+    obj = torch.rand(4, 4, 4, device="cuda")
+    # config.set lacks __exit__, so restore explicitly rather than via `with`.
+    config.set({"use_cuda_kernels": False})
+    try:
+        torch.testing.assert_close(tv_loss_vol_sq(obj), tv_vol_sq_ref(obj))
+    finally:
+        config.set({"use_cuda_kernels": True})
diff --git a/tests/tomography/test_dataset_models.py b/tests/tomography/test_dataset_models.py
index 04abd3b58..8aa9c27aa 100644
--- a/tests/tomography/test_dataset_models.py
+++ b/tests/tomography/test_dataset_models.py
@@ -8,6 +8,7 @@
 import pytest
 import torch
 
+from quantem.core.ml.optimizer_mixin import OptimizerParams
 from quantem.tomography.dataset_models import (
     DatasetConstraintParams,
     DatasetValue,
@@ -15,6 +16,7 @@
     TomographyINRPretrainDataset,
     TomographyPixDataset,
 )
+from quantem.tomography.tomography_opt import _pose_optimizer_specs
 
 from .conftest import requires_torch
 
@@ -58,6 +60,20 @@ def test_normalised_by_95th_quantile(self):
         q95 = torch.quantile(d.tilt_stack, 0.95)
         assert torch.isclose(q95, torch.tensor(1.0), atol=1e-4)
 
+    def test_sparse_stack_normalisation_stays_finite(self):
+        # >95% zeros makes the 95th quantile 0; normalisation must not divide
+        # by it (inf/NaN targets poison every parameter on the first backward).
+        stack = np.zeros((5, 12, 12), dtype=np.float32)
+        stack[:, 5:7, 5:7] = 10.0
+        d = TomographyPixDataset.from_data(stack, np.linspace(-60, 60, 5).astype(np.float32))
+        assert torch.isfinite(d.tilt_stack).all()
+        assert d.tilt_stack.max() > 0
+
+    def test_all_zero_stack_raises(self):
+        stack = np.zeros((5, 12, 12), dtype=np.float32)
+        with pytest.raises(ValueError):
+            TomographyPixDataset.from_data(stack, np.linspace(-60, 60, 5).astype(np.float32))
+
     def test_reference_idx_and_learnable_tilts(self):
         # negated angles -> [40, 15, -10, -35, -60]; smallest |angle| is index 2.
         angles = np.linspace(-40, 60, 5).astype(np.float32)
@@ -82,6 +98,83 @@ def test_to_materialises_pose_parameters(self):
 
 @requires_torch
 class TestTomographyINRDataset:
+    @pytest.mark.parametrize("shape", [(5,), (4, 11), (4, 12, 1)])
+    def test_tilt_angles_per_row_rejects_wrong_shape(self, shape):
+        stack = _stack(nang=4, n=12)
+        row_angles = np.zeros(shape, dtype=np.float32)
+
+        with pytest.raises(ValueError, match="tilt_angles_per_row must have shape"):
+            TomographyINRDataset.from_data(
+                stack,
+                np.linspace(-60, 60, 4, dtype="f4"),
+                tilt_angles_per_row=row_angles,
+            )
+
+    @pytest.mark.parametrize("shape", [(5,), (4, 11), (4, 12, 1)])
+    def test_tilt_angles_per_col_rejects_wrong_shape(self, shape):
+        stack = _stack(nang=4, n=12)
+        col_angles = np.zeros(shape, dtype=np.float32)
+
+        with pytest.raises(ValueError, match="tilt_angles_per_col must have shape"):
+            TomographyINRDataset.from_data(
+                stack,
+                np.linspace(-60, 60, 4, dtype="f4"),
+                tilt_angles_per_col=col_angles,
+            )
+
+    def test_per_row_and_per_col_angles_are_mutually_exclusive(self):
+        stack = _stack(nang=3, n=12)
+        pixel_angles = np.zeros(stack.shape[:2], dtype=np.float32)
+
+        with pytest.raises(ValueError, match="mutually exclusive"):
+            TomographyINRDataset.from_data(
+                stack,
+                np.linspace(-60, 60, 3, dtype="f4"),
+                tilt_angles_per_row=pixel_angles,
+                tilt_angles_per_col=pixel_angles,
+            )
+
+    def test_getitem_uses_legacy_frame_angle_without_per_row_angles(self):
+        angles = np.linspace(-60, 60, 5, dtype="f4")
+        d = TomographyINRDataset.from_data(_stack(nang=5, n=12), angles)
+
+        assert d.tilt_angles_per_row is None
+        assert d.tilt_angles_per_col is None
+        assert d[0]["phi"] == angles[0]
+        assert d[11 * 12]["phi"] == angles[0]
+
+    def test_getitem_uses_arbitrary_per_row_angles(self):
+        stack = _stack(nang=3, n=12)
+        row_angles = torch.linspace(-61.3, 47.9, stack.shape[0] * stack.shape[1]).reshape(
+            stack.shape[:2]
+        )
+        d = TomographyINRDataset.from_data(
+            stack,
+            np.linspace(-60, 60, 3, dtype="f4"),
+            tilt_angles_per_row=row_angles,
+        )
+
+        assert d.tilt_angles_per_row is row_angles
+        assert d[2 * 12]["phi"] == row_angles[0, 2]
+        assert d[9 * 12 + 7]["phi"] == row_angles[0, 9]
+        assert d[2 * 12]["phi"] != d[9 * 12 + 7]["phi"]
+
+    def test_getitem_uses_arbitrary_per_col_angles(self):
+        stack = _stack(nang=3, n=12)
+        col_angles = torch.linspace(-61.3, 47.9, stack.shape[0] * stack.shape[2]).reshape(
+            stack.shape[0], stack.shape[2]
+        )
+        d = TomographyINRDataset.from_data(
+            stack,
+            np.linspace(-60, 60, 3, dtype="f4"),
+            tilt_angles_per_col=col_angles,
+        )
+
+        assert d.tilt_angles_per_col is col_angles
+        assert d[2 * 12 + 3]["phi"] == col_angles[0, 3]
+        assert d[9 * 12 + 7]["phi"] == col_angles[0, 7]
+        assert d[2 * 12 + 3]["phi"] != d[9 * 12 + 7]["phi"]
+
     def test_len_is_projections_times_pixels(self):
         d = TomographyINRDataset.from_data(
             _stack(nang=5, n=12), np.linspace(-60, 60, 5, dtype="f4")
@@ -95,6 +188,98 @@ def test_getitem_keys(self):
         item = d[0]
         assert {"phi", "pixel_i", "pixel_j", "projection_idx", "target_value"} <= set(item.keys())
 
+    @pytest.mark.parametrize(
+        "learn_shift,learn_tilt_axis,shift_lr,tilt_axis_lr,expected_lrs",
+        [
+            (True, True, 2e-2, 3e-3, {"pose_shift": 2e-2, "pose_tilt_axis": 3e-3}),
+            (True, True, None, None, {"pose_shift": 1e-2, "pose_tilt_axis": 1e-2}),
+            (True, False, 2e-2, 3e-3, {"pose_shift": 2e-2}),
+            (False, True, 2e-2, 3e-3, {"pose_tilt_axis": 3e-3}),
+        ],
+    )
+    def test_pose_optimizer_groups_and_lrs(
+        self, learn_shift, learn_tilt_axis, shift_lr, tilt_axis_lr, expected_lrs
+    ):
+        angles = np.linspace(-60, 60, 5, dtype="f4")
+        d = TomographyINRDataset.from_data(
+            _stack(nang=len(angles), n=12),
+            angles,
+            learn_shift=learn_shift,
+            learn_tilt_axis=learn_tilt_axis,
+        )
+        d.to("cpu")
+        specs = _pose_optimizer_specs(
+            d, pose_lr=1e-2, shift_lr=shift_lr, tilt_axis_lr=tilt_axis_lr
+        )
+        d.set_optimizer(specs)
+
+        groups = {group["name"]: group for group in d.optimizer.param_groups}
+        assert set(groups) == set(expected_lrs)
+        for name, lr in expected_lrs.items():
+            assert groups[name]["lr"] == pytest.approx(lr)
+
+        if learn_shift:
+            shift_params = groups["pose_shift"]["params"]
+            assert len(shift_params) == 1 and shift_params[0] is d._shifts_params
+            assert d._shifts_params.shape == (len(angles) - 1, 2)
+        if learn_tilt_axis:
+            tilt_params = groups["pose_tilt_axis"]["params"]
+            assert len(tilt_params) == 2
+            assert tilt_params[0] is d._z1_params and tilt_params[1] is d._z3_params
+            assert d._z1_params.shape == d._z3_params.shape == (len(angles) - 1,)
+
+        optimized = [param for group in groups.values() for param in group["params"]]
+        assert all(
+            ref is not param
+            for ref in (d._shifts_ref, d._z1_ref, d._z3_ref)
+            for param in optimized
+        )
+
+    def test_legacy_shared_pose_optimizer_expands_to_active_groups(self):
+        d = TomographyINRDataset.from_data(
+            _stack(nang=5, n=12), np.linspace(-60, 60, 5, dtype="f4")
+        )
+        d.to("cpu")
+        d.set_optimizer(OptimizerParams.Adam(lr=1e-2))
+        assert {group["name"]: group["lr"] for group in d.optimizer.param_groups} == {
+            "pose_shift": 1e-2,
+            "pose_tilt_axis": 1e-2,
+        }
+
+    def test_full_tomography_pose_optimizer_groups_and_lrs(self):
+        from quantem.core.ml.inr import HSiren
+        from quantem.tomography.object_models import ObjectINR
+        from quantem.tomography.tomography import Tomography
+
+        model = HSiren(in_features=3, out_features=1, hidden_layers=1, hidden_features=8)
+        obj = ObjectINR.from_model(model, shape=(12, 12, 12), device="cpu")
+        d = TomographyINRDataset.from_data(
+            _stack(nang=5, n=12),
+            np.linspace(-60, 60, 5, dtype="f4"),
+            learn_shift=True,
+            learn_tilt_axis=True,
+        )
+        tomo = Tomography.from_models(dset=d, obj_model=obj, device="cpu", verbose=False)
+        tomo.optimizer_params = {
+            "pose": {
+                "pose_shift": {"name": "adam", "lr": 2e-2},
+                "pose_tilt_axis": {"name": "adam", "lr": 3e-3},
+            }
+        }
+        tomo.set_optimizers()
+
+        groups = {group["name"]: group for group in tomo.dset.optimizer.param_groups}
+        assert {name: group["lr"] for name, group in groups.items()} == {
+            "pose_shift": 2e-2,
+            "pose_tilt_axis": 3e-3,
+        }
+        optimized = [param for group in groups.values() for param in group["params"]]
+        assert all(
+            ref is not param
+            for ref in (tomo.dset._shifts_ref, tomo.dset._z1_ref, tomo.dset._z3_ref)
+            for param in optimized
+        )
+
     @pytest.mark.parametrize(
         "learn_shift,learn_tilt_axis",
         [(True, True), (True, False), (False, True), (False, False)],
@@ -161,11 +346,94 @@ def test_transform_batch_rays_identity_at_zero_pose(self):
         # No rotation and no shift -> rays pass through unchanged.
         assert torch.allclose(out, rays, atol=1e-5)
 
+    @pytest.mark.parametrize("ray_sampling", ["legacy", "box_fixed_ds"])
+    def test_per_row_angles_produce_rx_ray_directions(self, ray_sampling):
+        n = 5
+        stack = _stack(nang=2, n=n)
+        row_angles = np.array(
+            [[-37.2, -19.7, 3.4, 28.1, 52.6], [7.3, 8.8, 10.1, 11.9, 13.6]],
+            dtype=np.float32,
+        )
+        d = TomographyINRDataset.from_data(
+            stack,
+            np.array([0.0, 10.0], dtype=np.float32),
+            tilt_angles_per_row=row_angles,
+            ray_sampling=ray_sampling,
+        )
+        d.to("cpu")
+        items = [d[1 * n + 2], d[3 * n + 2]]
+        batch = {
+            key: torch.as_tensor([item[key] for item in items])
+            for key in ("projection_idx", "pixel_i", "pixel_j", "phi", "target_value")
+        }
+
+        coords = d.get_coords(batch, N=n, num_samples_per_ray=9)
+        if ray_sampling == "legacy":
+            sampled_rays = coords.reshape(2, 9, 3)
+        else:
+            ray_ids = d._ray_meta["local_ray_ids"]
+            sampled_rays = [coords[ray_ids == ray_idx] for ray_idx in range(2)]
+
+        actual_directions = torch.stack(
+            [
+                (ray[-1] - ray[0]) / torch.linalg.vector_norm(ray[-1] - ray[0])
+                for ray in sampled_rays
+            ]
+        )
+        rotations = d._compose_euler_rotation(torch.zeros(2), batch["phi"], torch.zeros(2))
+        expected_directions = rotations @ torch.tensor([0.0, 0.0, 1.0])
+
+        torch.testing.assert_close(actual_directions, expected_directions, atol=1e-6, rtol=1e-6)
+
+    @pytest.mark.parametrize("ray_sampling", ["legacy", "box_fixed_ds"])
+    def test_per_col_angles_produce_rx_ray_directions(self, ray_sampling):
+        n = 5
+        stack = _stack(nang=2, n=n)
+        col_angles = np.array(
+            [[-37.2, -19.7, 3.4, 28.1, 52.6], [7.3, 8.8, 10.1, 11.9, 13.6]],
+            dtype=np.float32,
+        )
+        d = TomographyINRDataset.from_data(
+            stack,
+            np.array([0.0, 10.0], dtype=np.float32),
+            tilt_angles_per_col=col_angles,
+            ray_sampling=ray_sampling,
+        )
+        d.to("cpu")
+        items = [d[2 * n + 1], d[2 * n + 3]]
+        batch = {
+            key: torch.as_tensor([item[key] for item in items])
+            for key in ("projection_idx", "pixel_i", "pixel_j", "phi", "target_value")
+        }
+
+        coords = d.get_coords(batch, N=n, num_samples_per_ray=9)
+        if ray_sampling == "legacy":
+            sampled_rays = coords.reshape(2, 9, 3)
+        else:
+            ray_ids = d._ray_meta["local_ray_ids"]
+            sampled_rays = [coords[ray_ids == ray_idx] for ray_idx in range(2)]
+
+        actual_directions = torch.stack(
+            [
+                (ray[-1] - ray[0]) / torch.linalg.vector_norm(ray[-1] - ray[0])
+                for ray in sampled_rays
+            ]
+        )
+        rotations = d._compose_euler_rotation(torch.zeros(2), batch["phi"], torch.zeros(2))
+        expected_directions = rotations @ torch.tensor([0.0, 0.0, 1.0])
+
+        torch.testing.assert_close(actual_directions, expected_directions, atol=1e-6, rtol=1e-6)
+
     def test_integrate_rays_sums_with_step_size(self):
+        # integrate_rays became an instance method dispatching on ray_sampling
+        # (box_fixed_ds vs legacy); the default legacy path keeps the original
+        # step-size summation this test pins down.
         B, S = 3, 5
-        out = TomographyINRDataset.integrate_rays(
-            torch.ones(B, S), num_samples_per_ray=S, target_values_len=B
+        stack = _stack(nang=3, n=4)
+        d = TomographyINRDataset.from_data(
+            stack, np.linspace(-60, 60, 3, dtype="f4"), ray_sampling="legacy"
         )
+        out = d.integrate_rays(torch.ones(B * S), num_samples_per_ray=S, target_values_len=B)
         step = 2.0 / (S - 1)
         assert out.shape == (B,)
         assert torch.allclose(out, torch.full((B,), S * step))
@@ -180,6 +448,25 @@ def test_getitem_index_mapping(self):
         assert int(item["pixel_j"]) == 1
         assert torch.isclose(item["target_value"], d.tilt_stack[1, 1, 1])
 
+    def test_len_and_getitem_non_square(self):
+        # Regression: pixel_i/pixel_j must decompose by the width (shape[2]) and __len__
+        # by H*W, not max(shape)^2 -- both were wrong for rectangular tilt images.
+        rng = np.random.default_rng(0)
+        stack = rng.random((3, 4, 6)).astype(np.float32)  # nang=3, H=4, W=6
+        d = TomographyINRDataset.from_data(stack, np.linspace(-60, 60, 3, dtype="f4"))
+        assert len(d) == 3 * 4 * 6
+        # idx = proj*(H*W) + i*W + j = 1*24 + 2*6 + 3 = 39
+        item = d[39]
+        assert int(item["projection_idx"]) == 1
+        assert int(item["pixel_i"]) == 2
+        assert int(item["pixel_j"]) == 3
+        assert torch.isclose(item["target_value"], d.tilt_stack[1, 2, 3])
+        # every index in range maps to a valid pixel
+        last = d[len(d) - 1]
+        assert int(last["projection_idx"]) == 2
+        assert int(last["pixel_i"]) == 3
+        assert int(last["pixel_j"]) == 5
+
     def test_save_load_parameters_roundtrip(self, tmp_path):
         angles = np.linspace(-60, 60, 5, dtype="f4")
         d = TomographyINRDataset.from_data(_stack(), angles)
@@ -192,3 +479,30 @@ def test_save_load_parameters_roundtrip(self, tmp_path):
         d2.to("cpu")
         d2.load_parameters(path)
         assert torch.allclose(d2.z1_params.detach(), d.z1_params.detach())
+
+
+@pytest.mark.parametrize("cls", [TomographyPixDataset, TomographyINRDataset])
+def test_to_preserves_trained_pose_parameters(cls):
+    """Regression: to() rebuilt the pose parameters from the initial-value buffers,
+    so any device move after training (e.g. from_file(...).to(device)) silently
+    reset the learned pose to zero."""
+    d = cls.from_data(_stack(), np.linspace(-60, 60, 5, dtype="f4"))
+    d.to("cpu")
+    d.z1_params.data.fill_(0.37)
+    d.z3_params.data.fill_(-0.21)
+    d.shifts_params.data.fill_(1.5)
+
+    d.to("cpu")
+
+    assert torch.allclose(d.z1_params.detach(), torch.full_like(d.z1_params, 0.37))
+    assert torch.allclose(d.z3_params.detach(), torch.full_like(d.z3_params, -0.21))
+    assert torch.allclose(d.shifts_params.detach(), torch.full_like(d.shifts_params, 1.5))
+
+
+def test_learnable_tilts_is_read_only():
+    """Regression: the old setter wrote a private attribute the getter never read,
+    so assignment appeared to succeed while silently doing nothing."""
+    d = TomographyPixDataset.from_data(_stack(), np.linspace(-60, 60, 5).astype(np.float32))
+    with pytest.raises(AttributeError):
+        d.learnable_tilts = 3
+    assert d.learnable_tilts == 4
diff --git a/tests/tomography/test_device_batch_sampler.py b/tests/tomography/test_device_batch_sampler.py
new file mode 100644
index 000000000..624885402
--- /dev/null
+++ b/tests/tomography/test_device_batch_sampler.py
@@ -0,0 +1,245 @@
+"""Tests for ``DeviceBatchSampler`` and its wiring in ``Tomography.reconstruct``.
+
+The sampler must yield batches identical in content to the per-pixel
+DataLoader path (same keys, same index decode as
+``TomographyINRDataset.__getitem__``), cover each pixel exactly once per
+epoch (minus the dropped tail batch), and respect the train/val index
+split. Device-independent, so everything here runs on CPU.
+"""
+
+import numpy as np
+import torch
+
+from quantem.tomography.dataset_models import DeviceBatchSampler, TomographyINRDataset
+
+
+def _dset(n_proj=4, n=10, seed=0):
+    rng = np.random.default_rng(seed)
+    stack = rng.random((n_proj, n, n)).astype(np.float32)
+    angles = np.linspace(-60, 60, n_proj).astype(np.float32)
+    return TomographyINRDataset.from_data(tilt_stack=stack, tilt_angles=angles)
+
+
+def _per_row_dset(n_proj=4, height=7, width=11, seed=0):
+    rng = np.random.default_rng(seed)
+    stack = rng.random((n_proj, height, width)).astype(np.float32)
+    angles = np.linspace(-60, 60, n_proj).astype(np.float32)
+    row_angles = np.linspace(-63.7, 58.9, n_proj * height, dtype=np.float32).reshape(
+        n_proj, height
+    )
+    return TomographyINRDataset.from_data(
+        tilt_stack=stack,
+        tilt_angles=angles,
+        tilt_angles_per_row=row_angles,
+    )
+
+
+def _per_col_dset(n_proj=4, height=7, width=11, seed=0):
+    rng = np.random.default_rng(seed)
+    stack = rng.random((n_proj, height, width)).astype(np.float32)
+    angles = np.linspace(-60, 60, n_proj).astype(np.float32)
+    col_angles = np.linspace(-63.7, 58.9, n_proj * width, dtype=np.float32).reshape(n_proj, width)
+    return TomographyINRDataset.from_data(
+        tilt_stack=stack,
+        tilt_angles=angles,
+        tilt_angles_per_col=col_angles,
+    )
+
+
+def _assert_batch_matches_getitem(dset, batch):
+    width = dset.tilt_stack.shape[2]
+    per_proj = dset.tilt_stack.shape[1] * width
+    flat_indices = batch["projection_idx"] * per_proj + batch["pixel_i"] * width + batch["pixel_j"]
+    for batch_idx, flat_idx in enumerate(flat_indices.tolist()):
+        item = dset[flat_idx]
+        for key in ("projection_idx", "pixel_i", "pixel_j", "phi", "target_value"):
+            torch.testing.assert_close(
+                batch[key][batch_idx],
+                torch.as_tensor(item[key], dtype=batch[key].dtype),
+                rtol=0,
+                atol=0,
+            )
+
+
+def test_batches_match_getitem():
+    dset = _dset()
+    sampler = DeviceBatchSampler(dset, batch_size=37, device="cpu", shuffle=False)
+    seen = 0
+    for batch in sampler:
+        for k in range(len(batch["target_value"])):
+            item = dset[seen + k]
+            # __getitem__ returns plain ints for the index keys (cheaper than
+            # 0-d tensors); coerce before comparing against the batched tensors.
+            for key in ("projection_idx", "pixel_i", "pixel_j", "phi", "target_value"):
+                torch.testing.assert_close(
+                    batch[key][k],
+                    torch.as_tensor(item[key], dtype=batch[key].dtype),
+                    rtol=0,
+                    atol=0,
+                )
+        seen += len(batch["target_value"])
+
+
+def test_non_square_per_row_batches_match_getitem_exactly():
+    dset = _per_row_dset()
+    sampler = DeviceBatchSampler(
+        dset,
+        batch_size=17,
+        device="cpu",
+        shuffle=False,
+        drop_last=False,
+    )
+
+    for batch in sampler:
+        _assert_batch_matches_getitem(dset, batch)
+
+
+def test_non_square_per_col_batches_match_getitem_exactly():
+    dset = _per_col_dset()
+    sampler = DeviceBatchSampler(
+        dset,
+        batch_size=17,
+        device="cpu",
+        shuffle=False,
+        drop_last=False,
+    )
+
+    for batch in sampler:
+        _assert_batch_matches_getitem(dset, batch)
+
+
+def test_per_row_angles_match_getitem_in_shuffled_train_and_val_split():
+    dset = _per_row_dset()
+    generator = torch.Generator().manual_seed(0)
+    perm = torch.randperm(len(dset), generator=generator)
+    n_val = len(dset) // 4
+    samplers = (
+        DeviceBatchSampler(
+            dset,
+            19,
+            "cpu",
+            indices=perm[n_val:],
+            shuffle=True,
+            drop_last=False,
+        ),
+        DeviceBatchSampler(
+            dset,
+            19,
+            "cpu",
+            indices=perm[:n_val],
+            shuffle=False,
+            drop_last=False,
+        ),
+    )
+
+    for sampler in samplers:
+        for batch in sampler:
+            _assert_batch_matches_getitem(dset, batch)
+
+
+def test_epoch_covers_indices_once_with_drop_last():
+    dset = _dset()
+    n = len(dset)
+    batch_size = 64
+    sampler = DeviceBatchSampler(dset, batch_size=batch_size, device="cpu", shuffle=True)
+    assert len(sampler) == n // batch_size
+    per_proj = dset.tilt_stack.shape[1] * dset.tilt_stack.shape[2]
+    flat = []
+    for batch in sampler:
+        assert len(batch["target_value"]) == batch_size
+        flat.append(
+            batch["projection_idx"] * per_proj
+            + batch["pixel_i"] * dset.tilt_stack.shape[1]
+            + batch["pixel_j"]
+        )
+    flat = torch.cat(flat)
+    assert flat.unique().numel() == flat.numel()  # no repeats within an epoch
+    assert flat.numel() == len(sampler) * batch_size
+
+
+def test_shuffle_changes_order_between_epochs():
+    dset = _dset()
+    sampler = DeviceBatchSampler(dset, batch_size=50, device="cpu", shuffle=True)
+    first = next(iter(sampler))["target_value"]
+    second = next(iter(sampler))["target_value"]
+    assert not torch.equal(first, second)
+
+
+def test_val_split_is_disjoint():
+    dset = _dset()
+    n = len(dset)
+    perm = torch.randperm(n)
+    n_val = n // 10
+    train = DeviceBatchSampler(dset, 32, "cpu", indices=perm[n_val:])
+    val = DeviceBatchSampler(dset, 32, "cpu", indices=perm[:n_val], shuffle=False)
+    assert len(train._indices) + len(val._indices) == n
+    assert torch.cat([train._indices, val._indices]).unique().numel() == n
+
+
+def test_reconstruct_uses_sampler_single_process():
+    """Smoke: the single-process reconstruct path builds DeviceBatchSamplers.
+
+    Object models that go through ``setup_distributed`` must be built on CUDA
+    when a CUDA device exists, so pick the device accordingly.
+    """
+    from quantem.core.ml.models.kplanes import KPlanesTILTED
+    from quantem.core.ml.optimizer_mixin import OptimizerParams
+    from quantem.tomography.object_models import ObjectINR
+    from quantem.tomography.tomography import Tomography
+
+    device = "cuda:0" if torch.cuda.is_available() else "cpu"
+    dset = _per_row_dset(n_proj=3, height=8, width=8)
+    model = KPlanesTILTED(M_features=2, resolution=(8, 8, 8), multiscale_res_multipliers=[1], T=1)
+    obj = ObjectINR.from_model(model, shape=(8, 8, 8), device=device)
+    tomo = Tomography.from_models(dset=dset, obj_model=obj, device=device, verbose=False)
+    tomo.reconstruct(
+        num_iter=1,
+        batch_size=32,
+        num_workers=0,
+        val_fraction=0.25,
+        optimizer_params={
+            "object": {"default": OptimizerParams.Adam(lr=1e-3)},
+            "pose": OptimizerParams.Adam(lr=1e-3),
+        },
+    )
+    assert isinstance(tomo.dataloader, DeviceBatchSampler)
+    assert isinstance(tomo.val_dataloader, DeviceBatchSampler)
+    # the loop drives epoch reshuffling through set_epoch on self.sampler
+    assert tomo.sampler is tomo.dataloader
+
+
+def _flat(batch, s1, s2):
+    return batch["projection_idx"] * (s1 * s2) + batch["pixel_i"] * s2 + batch["pixel_j"]
+
+
+def test_ddp_shards_are_disjoint_and_equal():
+    dset = _dset()
+    s1, s2 = dset.tilt_stack.shape[1], dset.tilt_stack.shape[2]
+    world = 4
+    shards = []
+    for rank in range(world):
+        sampler = DeviceBatchSampler(dset, 32, "cpu", rank=rank, world_size=world)
+        sampler.set_epoch(3)
+        flats = torch.cat([_flat(b, s1, s2) for b in sampler])
+        shards.append(flats)
+        assert len(sampler) == (len(dset) // world) // 32  # equal on every rank
+    allv = torch.cat(shards)
+    assert allv.unique().numel() == allv.numel()  # no pixel on two ranks
+
+
+def test_ddp_epoch_permutation_shared_and_reproducible():
+    dset = _dset()
+    s1, s2 = dset.tilt_stack.shape[1], dset.tilt_stack.shape[2]
+
+    def epoch_flats(rank, epoch):
+        sampler = DeviceBatchSampler(dset, 32, "cpu", rank=rank, world_size=2)
+        sampler.set_epoch(epoch)
+        return torch.cat([_flat(b, s1, s2) for b in sampler])
+
+    # same (rank, epoch) on a fresh instance -> identical batches
+    torch.testing.assert_close(epoch_flats(0, 7), epoch_flats(0, 7), rtol=0, atol=0)
+    # different epoch -> different order
+    assert not torch.equal(epoch_flats(0, 7), epoch_flats(0, 8))
+    # ranks of the same epoch are disjoint
+    both = torch.cat([epoch_flats(0, 7), epoch_flats(1, 7)])
+    assert both.unique().numel() == both.numel()
diff --git a/tests/tomography/test_grad_step_snapshots.py b/tests/tomography/test_grad_step_snapshots.py
new file mode 100644
index 000000000..4806e3a71
--- /dev/null
+++ b/tests/tomography/test_grad_step_snapshots.py
@@ -0,0 +1,124 @@
+"""Pure tests for tomography grad-step snapshot triggering and plumbing."""
+
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+from quantem.tomography import tomography
+
+
+class _ObjModel:
+    def __init__(self):
+        self.calls = 0
+        self.volume = np.arange(8, dtype=np.float64).reshape(1, 2, 2, 2)
+
+    @property
+    def obj_view(self):
+        self.calls += 1
+        return self.volume
+
+
+class _Logger:
+    def __init__(self):
+        self.scalars = []
+
+    def log_scalar(self, tag, value, step, step_domain="epoch"):
+        self.scalars.append((tag, value, step, step_domain))
+
+
+@pytest.mark.parametrize(
+    ("snapshot_every", "steps"),
+    [
+        (0, []),
+        (1, [1, 2, 3, 4, 5, 6]),
+        (2, [2, 4, 6]),
+        (3, [3, 6]),
+        (10, []),
+    ],
+)
+def test_should_take_grad_step_snapshot(snapshot_every, steps):
+    fired = [
+        step
+        for step in range(1, 7)
+        if tomography._should_take_grad_step_snapshot(step, snapshot_every)
+    ]
+    assert fired == steps
+
+
+def test_take_grad_step_snapshot_main_rank_writes_logs_and_callbacks(monkeypatch):
+    mkdir_calls = []
+    saved = {}
+
+    def fake_mkdir(self, parents=False, exist_ok=False):
+        mkdir_calls.append((self, parents, exist_ok))
+
+    def fake_save(path, array):
+        saved["path"] = path
+        saved["array"] = array
+
+    monkeypatch.setattr(tomography.Path, "mkdir", fake_mkdir)
+    monkeypatch.setattr(tomography.np, "save", fake_save)
+
+    obj_model = _ObjModel()
+    logger = _Logger()
+    callback_calls = []
+
+    tomography._take_grad_step_snapshot(
+        obj_model=obj_model,
+        grad_step=4,
+        global_rank=0,
+        logger=logger,
+        snapshot_dir="snapshots",
+        snapshot_callback=lambda step, volume: callback_calls.append((step, volume)),
+    )
+
+    assert obj_model.calls == 1
+    assert mkdir_calls == [(Path("snapshots"), True, True)]
+    assert saved["path"] == Path("snapshots") / "step_4.npy"
+    assert saved["array"].dtype == np.float32
+    np.testing.assert_array_equal(saved["array"], obj_model.volume.astype(np.float32))
+    assert logger.scalars == [("snapshots/last_grad_step", 4.0, 4, "grad_step")]
+    assert callback_calls == [(4, obj_model.volume)]
+
+
+def test_take_grad_step_snapshot_non_main_rank_gets_collective_but_no_volume(monkeypatch):
+    monkeypatch.setattr(
+        tomography.np,
+        "save",
+        lambda *args, **kwargs: pytest.fail("non-main ranks must not write snapshots"),
+    )
+
+    obj_model = _ObjModel()
+    logger = _Logger()
+    callback_calls = []
+
+    tomography._take_grad_step_snapshot(
+        obj_model=obj_model,
+        grad_step=6,
+        global_rank=1,
+        logger=logger,
+        snapshot_dir="snapshots",
+        snapshot_callback=lambda step, volume: callback_calls.append((step, volume)),
+    )
+
+    assert obj_model.calls == 1
+    assert logger.scalars == []
+    assert callback_calls == [(6, None)]
+
+
+def test_take_grad_step_snapshot_callback_exception_propagates():
+    obj_model = _ObjModel()
+
+    def raise_callback(step, volume):
+        raise RuntimeError(f"stop at {step}")
+
+    with pytest.raises(RuntimeError, match="stop at 8"):
+        tomography._take_grad_step_snapshot(
+            obj_model=obj_model,
+            grad_step=8,
+            global_rank=1,
+            logger=None,
+            snapshot_dir=None,
+            snapshot_callback=raise_callback,
+        )
diff --git a/tests/tomography/test_kplanes_ms_tv_reconstruction.py b/tests/tomography/test_kplanes_ms_tv_reconstruction.py
new file mode 100644
index 000000000..5d35a0d01
--- /dev/null
+++ b/tests/tomography/test_kplanes_ms_tv_reconstruction.py
@@ -0,0 +1,127 @@
+"""CPU integration coverage for reconstruction-scoped multiscale plane-TV fusion."""
+
+import sys
+from types import ModuleType
+
+import pytest
+import torch
+
+from quantem.core import config
+from quantem.core.ml.models.kplanes import KPlanesTILTED
+from quantem.tomography.object_models import ObjectTensorDecomp
+from quantem.tomography.tomography_context import ReconstructionContext
+
+
+class _PretendCudaTensor(torch.Tensor):
+    """CPU tensor subclass that reaches the mocked optional-CUDA dispatch seam."""
+
+    @property
+    def is_cuda(self):
+        return True
+
+
+@pytest.fixture(autouse=True)
+def preserve_torch_rng_state():
+    state = torch.random.get_rng_state()
+    yield
+    torch.random.set_rng_state(state)
+
+
+@pytest.fixture
+def mocked_cuda_ml(monkeypatch):
+    calls = {"ms": 0, "ms_tv": 0}
+    cuda_ml = ModuleType("quantem.cuda.core.ml")
+
+    def multiscale(pts, rotations, grid0, grid1, grid2, *gates):
+        calls["ms"] += 1
+        width = rotations.shape[0] * sum(grid.shape[1] for grid in (grid0, grid1, grid2))
+        return torch.zeros((pts.shape[0], width), dtype=grid0.dtype)
+
+    def multiscale_tv(pts, rotations, grid0, grid1, grid2, *gates):
+        calls["ms_tv"] += 1
+        features = multiscale(pts, rotations, grid0, grid1, grid2, *gates)
+        calls["ms"] -= 1
+        return features, torch.tensor(2.5)
+
+    cuda_ml.kplanes_tilted_fuse = lambda pts, rotations, grid: None
+    cuda_ml._kplanes_tilted_fuse_builtin = cuda_ml.kplanes_tilted_fuse
+    cuda_ml.kplanes_tilted_fuse_ms = multiscale
+    cuda_ml.kplanes_tilted_fuse_ms_tv = multiscale_tv
+
+    cuda = ModuleType("quantem.cuda")
+    cuda.__path__ = []
+    cuda_core = ModuleType("quantem.cuda.core")
+    cuda_core.__path__ = []
+    cuda.core = cuda_core
+    cuda_core.ml = cuda_ml
+    monkeypatch.setattr(sys.modules["quantem"], "cuda", cuda, raising=False)
+    monkeypatch.setitem(sys.modules, "quantem.cuda", cuda)
+    monkeypatch.setitem(sys.modules, "quantem.cuda.core", cuda_core)
+    monkeypatch.setitem(sys.modules, "quantem.cuda.core.ml", cuda_ml)
+    monkeypatch.setattr(config, "get", lambda key, default=None: True)
+    monkeypatch.delenv("QUANTEM_KPLANES_MS_TV_FUSED", raising=False)
+    return calls
+
+
+@pytest.mark.skipif(
+    torch.cuda.is_available(),
+    reason=(
+        "CPU-only mock integration test; real GPU path covered by "
+        "test_kplanes_ms_tv_optimizer.py and the three-level cuda_graphs arm"
+    ),
+)
+def test_reconstruction_scope_requests_and_consumes_fused_plane_tv_once(
+    mocked_cuda_ml, monkeypatch
+):
+    model = KPlanesTILTED(
+        T=1,
+        M_features=2,
+        resolution=(4, 4, 4),
+        multiscale_res_multipliers=(1, 2, 3),
+    )
+    obj = ObjectTensorDecomp.from_model(model, shape=(4, 4, 4), device="cpu")
+    obj.constraints.tv_plane = 0.2
+    coords = torch.zeros((3, 3)).as_subclass(_PretendCudaTensor)
+
+    generic_density = obj.forward(coords)
+    assert generic_density.shape == (3,)
+    assert mocked_cuda_ml == {"ms": 1, "ms_tv": 0}
+
+    with obj.reconstruction_forward_context():
+        reconstruction_density = obj.forward(coords)
+    assert reconstruction_density.shape == (3,)
+    assert mocked_cuda_ml == {"ms": 1, "ms_tv": 1}
+    assert model._plane_tv_fusion_requested is False
+
+    def standalone_path_must_not_run():
+        raise AssertionError("standalone plane TV was computed after the fused auxiliary")
+
+    monkeypatch.setattr(obj, "_get_plane_tv_loss", standalone_path_must_not_run)
+    loss = obj.get_tv_loss(ReconstructionContext(coords=coords, pred=reconstruction_density))
+    torch.testing.assert_close(loss, torch.tensor(0.5))
+    assert obj._fused_plane_tv_loss is None
+
+
+@pytest.mark.skipif(
+    torch.cuda.is_available(),
+    reason=(
+        "CPU-only mock integration test; real GPU path covered by "
+        "test_kplanes_ms_tv_optimizer.py and the three-level cuda_graphs arm"
+    ),
+)
+def test_reconstruction_scope_clears_fused_plane_tv_after_exception(mocked_cuda_ml):
+    model = KPlanesTILTED(
+        T=1,
+        M_features=2,
+        resolution=(4, 4, 4),
+        multiscale_res_multipliers=(1, 2, 3),
+    )
+    obj = ObjectTensorDecomp.from_model(model, shape=(4, 4, 4), device="cpu")
+
+    with pytest.raises(RuntimeError, match="failed reconstruction"):
+        with obj.reconstruction_forward_context():
+            obj._fused_plane_tv_loss = torch.tensor(2.5, requires_grad=True)
+            raise RuntimeError("failed reconstruction")
+
+    assert obj._fused_plane_tv_loss is None
+    assert model._plane_tv_fusion_requested is False
diff --git a/tests/tomography/test_logger_tomography.py b/tests/tomography/test_logger_tomography.py
index 88d6bb265..4dab1a597 100644
--- a/tests/tomography/test_logger_tomography.py
+++ b/tests/tomography/test_logger_tomography.py
@@ -8,9 +8,13 @@
 headless.
 """
 
+import os
+import sys
 from types import SimpleNamespace
 
+import matplotlib.pyplot as plt
 import numpy as np
+import pytest
 import torch
 
 from quantem.tomography.logger_tomography import LoggerTomography
@@ -101,3 +105,130 @@ def test_log_iter_images(tmp_path):
         assert list(logger.log_dir.glob("events.out.tfevents.*"))
     finally:
         logger.close()
+
+
+def test_invalid_mode_raises(tmp_path):
+    with pytest.raises(ValueError, match="tensorboard.*wandb"):
+        LoggerTomography(
+            log_dir=str(tmp_path),
+            run_prefix="test_tomo",
+            mode="invalid",
+        )
+
+
+def test_wandb_mode_logs_under_run_dir(tmp_path, monkeypatch):
+    monkeypatch.delenv("WANDB_MODE", raising=False)
+    init_calls = []
+    logged = []
+    defined_metrics = []
+
+    class FakeConfig(dict):
+        def update(self, values, allow_val_change=False):
+            super().update(values)
+
+    class FakeRun:
+        def __init__(self):
+            self.config = FakeConfig()
+            self.finished = False
+
+        def define_metric(self, name, **kwargs):
+            defined_metrics.append((name, kwargs))
+
+        def log(self, data, **kwargs):
+            logged.append((data, kwargs))
+
+        def finish(self):
+            self.finished = True
+
+    def fake_init(**kwargs):
+        init_calls.append(kwargs)
+        return FakeRun()
+
+    fake_wandb = SimpleNamespace(
+        Image=lambda image: ("image", image),
+        Histogram=lambda values: ("histogram", values),
+        init=fake_init,
+    )
+    monkeypatch.setitem(sys.modules, "wandb", fake_wandb)
+
+    logger = LoggerTomography(
+        log_dir=str(tmp_path),
+        run_prefix="test_tomo",
+        run_suffix="wandb",
+        log_images_every=1,
+        mode="wandb",
+        wandb_config={"batch_size": 4},
+    )
+    try:
+        assert logger.mode == "wandb"
+        assert os.environ["WANDB_MODE"] == "offline"
+        assert logger.log_dir.exists()
+        assert (logger.log_dir / "wandb").exists()
+        assert init_calls[-1]["dir"] == str(logger.log_dir)
+        assert init_calls[-1]["config"] == {"batch_size": 4}
+
+        logger.attach_config({"num_iter": 2})
+        extra_steps = {"grad_step": 10}
+        logger.log_scalar("loss/total", 1.0, 1, extra_steps=extra_steps)
+        logger.log_image(
+            "volume/sum_z_0", np.ones((2, 2), dtype=np.float32), 1, extra_steps=extra_steps
+        )
+        logger.log_histogram(
+            "weights/object",
+            np.array([0.0, 1.0], dtype=np.float32),
+            1,
+            extra_steps=extra_steps,
+        )
+        logger.log_text("config/notes", "offline test", 1, extra_steps=extra_steps)
+        fig, ax = plt.subplots()
+        ax.plot([0, 1], [1, 0])
+        logger.log_figure("figures/test", fig, 1, extra_steps=extra_steps)
+        plt.close(fig)
+        logger.log_scalar(
+            "snapshots/last_grad_step", 48.0, 48, step_domain="grad_step"
+        )
+        logger.log_epoch(
+            epoch=2,
+            loss=0.8,
+            tilt_series_loss=0.6,
+            soft_loss=0.2,
+            grad_step=20,
+        )
+        logger.log_iter(
+            object_model=SimpleNamespace(_soft_constraint_losses=[0.1]),
+            iter=3,
+            consistency_loss=0.4,
+            total_loss=0.5,
+            learning_rates={"object": 1e-3},
+            num_samples_per_ray=8,
+            val_loss=0.3,
+            val_fg_loss=0.2,
+            val_bg_loss=0.1,
+            grad_step=30,
+        )
+        logger.flush()
+    finally:
+        logger.close()
+
+    assert defined_metrics == [
+        ("*", {"step_metric": "epoch"}),
+        ("snapshots/*", {"step_metric": "grad_step"}),
+    ]
+    assert all("step" not in kwargs for _, kwargs in logged)
+    assert logged[0] == ({"loss/total": 1.0, "grad_step": 10, "epoch": 1}, {})
+    assert "volume/sum_z_0" in logged[1][0]
+    assert logged[1][0]["epoch"] == 1
+    assert logged[1][0]["grad_step"] == 10
+    assert logged[2][0]["epoch"] == 1
+    assert logged[2][0]["grad_step"] == 10
+    assert logged[3] == ({"config/notes": "offline test", "grad_step": 10, "epoch": 1}, {})
+    assert "figures/test" in logged[4][0]
+    assert logged[4][0]["epoch"] == 1
+    assert logged[4][0]["grad_step"] == 10
+    assert logged[5] == ({"snapshots/last_grad_step": 48.0, "grad_step": 48}, {})
+    assert logged[6] == ({"loss/total": 0.8, "grad_step": 20, "epoch": 2}, {})
+    assert logged[7] == ({"loss/tilt_series": 0.6, "grad_step": 20, "epoch": 2}, {})
+    assert logged[8] == ({"loss/soft": 0.2, "grad_step": 20, "epoch": 2}, {})
+    for data, _ in logged[9:]:
+        assert data["epoch"] == 3
+        assert data["grad_step"] == 30
diff --git a/tests/tomography/test_object_models.py b/tests/tomography/test_object_models.py
index c4216827d..c0327d516 100644
--- a/tests/tomography/test_object_models.py
+++ b/tests/tomography/test_object_models.py
@@ -9,6 +9,7 @@
 import pytest
 import torch
 
+from quantem.core.ml.activation_functions import TruncExpActivation
 from quantem.tomography.object_models import (
     ObjConstraintParams,
     ObjectBase,
@@ -148,6 +149,17 @@ def test_soft_constraint_zero_when_tv_off(self):
         obj.constraints.tv_vol = 0.0
         assert float(obj.apply_soft_constraints(ctx).detach()) == 0.0
 
+    def test_soft_constraint_with_tv_on_backprops(self):
+        # Regression: soft_loss was a leaf tensor with requires_grad=True, so the
+        # in-place `+= tv_loss` raised RuntimeError whenever tv_vol > 0.
+        ctx = ReconstructionContext(obj=torch.rand(8, 8, 8).requires_grad_(True))
+        obj = ObjectPixelated.from_uniform(shape=(8, 8, 8), device="cpu")
+        obj.constraints.tv_vol = 1.0
+        loss = obj.apply_soft_constraints(ctx)
+        assert torch.isfinite(loss) and loss > 0.0
+        loss.backward()
+        assert ctx.obj.grad is not None
+
 
 class TestFactoryGuard:
     def test_objectbase_requires_token(self):
@@ -205,10 +217,25 @@ def _obj(self, device, n=16):
 
     def test_forward_masks_out_of_range(self, torch_device):
         obj = self._obj(torch_device)
-        coords = torch.tensor([[0.0, 0.0, 0.0], [5.0, 0.0, 0.0]], device=torch_device)
+        # Regression: the mask used to check x and y only, so tilted rays leaving the
+        # volume along z still contributed (extrapolated) density to the integral.
+        coords = torch.tensor(
+            [[0.0, 0.0, 0.0], [5.0, 0.0, 0.0], [0.0, 5.0, 0.0], [0.0, 0.0, 5.0]],
+            device=torch_device,
+        )
         out = obj.forward(coords)
-        assert out.shape[0] == 2
-        assert float(out[1].detach()) == 0.0  # x=5 is outside [-1, 1] -> masked to zero
+        assert out.shape[0] == 4
+        for i in (1, 2, 3):  # each axis outside [-1, 1] -> masked to zero
+            assert float(out[i].detach()) == 0.0
+
+    def test_soft_constraints_without_coords(self, torch_device):
+        # Regression: soft_loss was created on ctx.coords.device before the None
+        # check, raising AttributeError when no constraints are active.
+        obj = self._obj(torch_device)
+        obj.constraints.tv_vol = 0.0
+        obj.constraints.sparsity = 0.0
+        loss = obj.apply_soft_constraints(ReconstructionContext())
+        assert float(loss.detach()) == 0.0
 
     def test_apply_hard_constraints_positivity(self, torch_device):
         obj = self._obj(torch_device)
@@ -242,6 +269,30 @@ def test_apply_hard_constraints_positivity(self, torch_device):
         pred = torch.tensor([-2.0, 0.0, 3.0], device=torch_device)
         assert torch.all(obj.apply_hard_constraints(pred) >= 0.0)
 
+    def test_nonnegative_activation_capability_bypasses_redundant_clamp(
+        self, torch_device, monkeypatch
+    ):
+        obj = self._obj(torch_device)
+        obj.constraints.positivity = True
+        _model = obj.model.module if hasattr(obj.model, "module") else obj.model
+        _model.density_activation = TruncExpActivation(offset=1.0)
+        pred = torch.tensor([-2.0, 0.5], device=torch_device)
+
+        monkeypatch.delenv("QUANTEM_DENSITY_TAIL_FUSED", raising=False)
+        assert torch.equal(obj.apply_hard_constraints(pred), pred)
+
+        monkeypatch.setenv("QUANTEM_DENSITY_TAIL_FUSED", "0")
+        assert torch.equal(obj.apply_hard_constraints(pred), pred.clamp(min=0.0))
+
+    def test_unmarked_activation_keeps_positivity_clamp(self, torch_device, monkeypatch):
+        obj = self._obj(torch_device)
+        obj.constraints.positivity = True
+        _model = obj.model.module if hasattr(obj.model, "module") else obj.model
+        _model.density_activation = lambda values: torch.exp(values)
+        monkeypatch.delenv("QUANTEM_DENSITY_TAIL_FUSED", raising=False)
+        pred = torch.tensor([-2.0, 0.5], device=torch_device)
+        assert torch.equal(obj.apply_hard_constraints(pred), pred.clamp(min=0.0))
+
     def test_plane_tv_loss_nonneg_scalar(self, torch_device):
         obj = self._obj(torch_device)
         obj.constraints.tv_plane = 0.1
@@ -257,6 +308,39 @@ def test_volume_tv_loss_scalar(self, torch_device):
         assert loss.ndim == 0
         assert torch.isfinite(loss)
 
+    def test_soft_constraints_plane_tv_only(self, torch_device):
+        """Regression: tv_plane > 0 with tv_vol = 0 must still apply the plane TV.
+        The gate previously keyed on tv_vol alone, silently dropping plane-only TV."""
+        obj = self._obj(torch_device)
+        obj.constraints.tv_plane = 0.1
+        obj.constraints.tv_vol = 0.0
+        coords = torch.rand(64, 3, device=torch_device) * 2 - 1
+        ctx = ReconstructionContext(coords=coords, pred=torch.zeros(64, device=torch_device))
+        loss = obj.apply_soft_constraints(ctx)
+        assert float(loss.detach()) > 0.0
+
+    def test_combined_plane_tv_value_and_weight_are_consumed_exactly_once(
+        self, torch_device, monkeypatch
+    ):
+        obj = self._obj(torch_device)
+        obj.constraints.tv_plane = 0.125
+        obj.constraints.tv_vol = 0.0
+        raw_tv = torch.tensor(4.0, device=torch_device, requires_grad=True)
+        obj._fused_plane_tv_loss = raw_tv
+
+        def separate_path_must_not_run():
+            raise AssertionError("combined TV was double-counted through the separate path")
+
+        monkeypatch.setattr(obj, "_get_plane_tv_loss", separate_path_must_not_run)
+        coords = torch.zeros(1, 3, device=torch_device)
+        ctx = ReconstructionContext(coords=coords, pred=torch.zeros(1, device=torch_device))
+        loss = obj.get_tv_loss(ctx)
+
+        torch.testing.assert_close(loss, torch.tensor(0.5, device=torch_device))
+        loss.backward()
+        torch.testing.assert_close(raw_tv.grad, torch.tensor(0.125, device=torch_device))
+        assert obj._fused_plane_tv_loss is None
+
     def test_normalize_optimizer_params_rejects_non_dict(self, torch_device):
         from quantem.core.ml.optimizer_mixin import OptimizerParams
 
@@ -272,3 +356,28 @@ def test_normalize_optimizer_params_rejects_wrong_keys(self, torch_device):
             obj._normalize_optimizer_params(
                 {"grids": OptimizerParams.Adam(), "wrong": OptimizerParams.Adam()}
             )
+
+
+class TestObjectINRCompileOptIn:
+    """compile_model=True must change performance only, never results."""
+
+    def _make(self, compile_model, device):
+        from quantem.core.ml.inr import HSiren
+
+        torch.manual_seed(0)
+        model = HSiren(hidden_layers=1, hidden_features=16, alpha=1)
+        return ObjectINR.from_model(
+            model, shape=(8, 8, 8), device=device, compile_model=compile_model
+        )
+
+    def test_default_is_eager(self, torch_device):
+        obj = self._make(False, torch_device)
+        assert obj._compile_model is False
+
+    @pytest.mark.slow
+    def test_compiled_forward_matches_eager(self, torch_device):
+        torch.manual_seed(1)
+        coords = (torch.rand(64, 3) * 2 - 1).to(torch_device)
+        out_eager = self._make(False, torch_device).forward(coords)
+        out_compiled = self._make(True, torch_device).forward(coords)
+        torch.testing.assert_close(out_compiled, out_eager, rtol=1e-5, atol=1e-6)
diff --git a/tests/tomography/test_pixel_holdout.py b/tests/tomography/test_pixel_holdout.py
new file mode 100644
index 000000000..be05d6643
--- /dev/null
+++ b/tests/tomography/test_pixel_holdout.py
@@ -0,0 +1,139 @@
+import numpy as np
+import torch
+
+from quantem.tomography.dataset_models import (
+    DeviceBatchSampler,
+    TomographyINRDataset,
+    build_pixel_holdout_split,
+)
+
+
+def _stack(n_proj=3, n=10):
+    stack = np.zeros((n_proj, n, n), dtype=np.float32)
+    foreground = stack.reshape(-1)[::7]
+    values = np.linspace(1.0, 5.0, min(20, len(foreground)), dtype=np.float32)
+    foreground[: len(values)] = values
+    return stack
+
+
+def _flat(batch, n):
+    return batch["projection_idx"] * (n * n) + batch["pixel_i"] * n + batch["pixel_j"]
+
+
+def test_holdout_split_is_seeded_and_foreground_aware():
+    stack = _stack()
+    first = build_pixel_holdout_split(stack, holdout_fraction=0.2, holdout_seed=11)
+    repeat = build_pixel_holdout_split(stack, holdout_fraction=0.2, holdout_seed=11)
+    different = build_pixel_holdout_split(stack, holdout_fraction=0.2, holdout_seed=12)
+
+    torch.testing.assert_close(first.val_indices, repeat.val_indices, rtol=0, atol=0)
+    assert not torch.equal(first.val_indices, different.val_indices)
+    assert first.val_fg_indices.numel() > 0
+    assert first.val_bg_indices.numel() > 0
+
+
+def test_holdout_fraction_arithmetic_is_exact():
+    stack = _stack(n_proj=2, n=10)
+    split = build_pixel_holdout_split(stack, holdout_fraction=0.15, holdout_seed=0)
+    n_pixels = stack.size
+    n_holdout = int(n_pixels * 0.15)
+
+    assert split.val_indices.numel() == n_holdout
+    assert split.train_indices.numel() == n_pixels - n_holdout
+    assert split.val_fg_indices.numel() + split.val_bg_indices.numel() == n_holdout
+
+
+def test_training_and_holdout_indices_are_disjoint():
+    stack = _stack()
+    split = build_pixel_holdout_split(stack, holdout_fraction=0.25, holdout_seed=3)
+    all_indices = torch.cat([split.train_indices, split.val_indices])
+
+    assert all_indices.unique().numel() == stack.size
+    assert torch.isin(split.train_indices, split.val_indices).sum().item() == 0
+
+
+def test_ddp_training_shards_exclude_same_holdout_pixels():
+    stack = _stack(n_proj=4, n=12)
+    angles = np.linspace(-60, 60, stack.shape[0], dtype=np.float32)
+    dset = TomographyINRDataset.from_data(stack, angles)
+    split = build_pixel_holdout_split(dset.tilt_stack, holdout_fraction=0.1, holdout_seed=9)
+    holdout = set(split.val_indices.tolist())
+
+    shards = []
+    for rank in range(3):
+        sampler = DeviceBatchSampler(
+            dset,
+            batch_size=16,
+            device="cpu",
+            indices=split.train_indices,
+            rank=rank,
+            world_size=3,
+        )
+        sampler.set_epoch(0)
+        flats = torch.cat([_flat(batch, stack.shape[1]) for batch in sampler])
+        assert not holdout.intersection(flats.tolist())
+        shards.append(flats)
+
+    all_train_seen = torch.cat(shards)
+    assert all_train_seen.unique().numel() == all_train_seen.numel()
+
+
+def test_validation_sampler_keeps_partial_final_batch():
+    stack = _stack(n_proj=2, n=8)
+    angles = np.linspace(-60, 60, stack.shape[0], dtype=np.float32)
+    dset = TomographyINRDataset.from_data(stack, angles)
+    split = build_pixel_holdout_split(stack, holdout_fraction=0.05, holdout_seed=4)
+
+    sampler = DeviceBatchSampler(
+        dset,
+        batch_size=1024,
+        device="cpu",
+        indices=split.val_indices,
+        shuffle=False,
+        drop_last=False,
+    )
+
+    assert len(sampler) == 1
+    batch = next(iter(sampler))
+    assert len(batch["target_value"]) == split.val_indices.numel()
+
+
+def test_per_row_angles_apply_to_all_holdout_samplers():
+    stack = _stack(n_proj=4, n=12)
+    angles = np.linspace(-60, 60, stack.shape[0], dtype=np.float32)
+    row_angles = np.linspace(
+        -62.4, 59.1, stack.shape[0] * stack.shape[1], dtype=np.float32
+    ).reshape(stack.shape[:2])
+    dset = TomographyINRDataset.from_data(
+        stack,
+        angles,
+        tilt_angles_per_row=row_angles,
+    )
+    split = build_pixel_holdout_split(dset.tilt_stack, holdout_fraction=0.2, holdout_seed=7)
+    sampler_configs = (
+        (split.train_indices, True),
+        (split.val_indices, False),
+        (split.val_fg_indices, False),
+        (split.val_bg_indices, False),
+    )
+
+    for indices, shuffle in sampler_configs:
+        sampler = DeviceBatchSampler(
+            dset,
+            batch_size=17,
+            device="cpu",
+            indices=indices,
+            shuffle=shuffle,
+            drop_last=False,
+        )
+        for batch in sampler:
+            flat_indices = _flat(batch, stack.shape[1])
+            for batch_idx, flat_idx in enumerate(flat_indices.tolist()):
+                item = dset[flat_idx]
+                for key in ("projection_idx", "pixel_i", "pixel_j", "phi", "target_value"):
+                    torch.testing.assert_close(
+                        batch[key][batch_idx],
+                        torch.as_tensor(item[key], dtype=batch[key].dtype),
+                        rtol=0,
+                        atol=0,
+                    )
diff --git a/tests/tomography/test_plane_tv_cuda_dispatch.py b/tests/tomography/test_plane_tv_cuda_dispatch.py
new file mode 100644
index 000000000..3a601989c
--- /dev/null
+++ b/tests/tomography/test_plane_tv_cuda_dispatch.py
@@ -0,0 +1,103 @@
+"""CPU-only capability and kill-switch coverage for fused plane TV."""
+
+import sys
+from types import ModuleType, SimpleNamespace
+
+import pytest
+import torch
+
+from quantem.core import config
+from quantem.tomography import object_models
+
+
+def _eager_reference(grids, rotations):
+    levels = []
+    for grid in grids:
+        dh = (grid[:, :, 1:, :] - grid[:, :, :-1, :]).pow(2).mean(dim=(1, 2, 3))
+        dw = (grid[:, :, :, 1:] - grid[:, :, :, :-1]).pow(2).mean(dim=(1, 2, 3))
+        levels.append((dh + dw).view(rotations, 3).sum(dim=1).mean())
+    return torch.stack(levels).sum()
+
+
+@pytest.fixture
+def mocked_plane_tv(monkeypatch):
+    calls = []
+    cuda_ml = ModuleType("quantem.cuda.core.ml")
+
+    def fused(*grids):
+        calls.append(grids)
+        return torch.tensor(7.0)
+
+    cuda_ml.plane_tv_loss = fused
+    cuda = ModuleType("quantem.cuda")
+    cuda.__path__ = []
+    cuda_core = ModuleType("quantem.cuda.core")
+    cuda_core.__path__ = []
+    cuda.core = cuda_core
+    cuda_core.ml = cuda_ml
+    monkeypatch.setattr(sys.modules["quantem"], "cuda", cuda, raising=False)
+    monkeypatch.setitem(sys.modules, "quantem.cuda", cuda)
+    monkeypatch.setitem(sys.modules, "quantem.cuda.core", cuda_core)
+    monkeypatch.setitem(sys.modules, "quantem.cuda.core.ml", cuda_ml)
+    monkeypatch.setattr(config, "get", lambda key, default=None: True)
+    monkeypatch.delenv("QUANTEM_PLANE_TV_FUSED", raising=False)
+    return cuda_ml, calls
+
+
+def _pretend_cuda_grids():
+    return tuple(SimpleNamespace(is_cuda=True, dtype=torch.float32) for _ in range(3))
+
+
+def test_three_cuda_levels_dispatch_to_fused_op(mocked_plane_tv):
+    _, calls = mocked_plane_tv
+    grids = _pretend_cuda_grids()
+    loss = object_models._plane_tv_loss(grids, tilted=True, rotations=2)
+    assert loss.item() == 7.0
+    assert len(calls) == 1
+    assert all(actual is expected for actual, expected in zip(calls[0], grids))
+
+
+def test_cpu_levels_use_eager_fallback(mocked_plane_tv):
+    _, calls = mocked_plane_tv
+    generator = torch.Generator().manual_seed(1)
+    grids = tuple(
+        torch.rand((6, channels, height, width), generator=generator)
+        for channels, height, width in ((2, 5, 7), (3, 9, 11), (4, 13, 17))
+    )
+    actual = object_models._plane_tv_loss(grids, tilted=True, rotations=2)
+    torch.testing.assert_close(actual, _eager_reference(grids, rotations=2))
+    assert calls == []
+
+
+def test_missing_cuda_capability_uses_eager_fallback(mocked_plane_tv, monkeypatch):
+    cuda_ml, calls = mocked_plane_tv
+    del cuda_ml.plane_tv_loss
+    expected = torch.tensor(3.5)
+    eager_calls = []
+
+    def eager(grids, tilted, rotations):
+        eager_calls.append((grids, tilted, rotations))
+        return expected
+
+    monkeypatch.setattr(object_models, "_plane_tv_loss_eager", eager)
+    actual = object_models._plane_tv_loss(_pretend_cuda_grids(), tilted=True, rotations=2)
+    assert actual is expected
+    assert len(eager_calls) == 1
+    assert calls == []
+
+
+def test_env_kill_switch_forces_eager_fallback(mocked_plane_tv, monkeypatch):
+    _, calls = mocked_plane_tv
+    monkeypatch.setenv("QUANTEM_PLANE_TV_FUSED", "0")
+    expected = torch.tensor(2.25)
+    eager_calls = []
+
+    def eager(grids, tilted, rotations):
+        eager_calls.append((grids, tilted, rotations))
+        return expected
+
+    monkeypatch.setattr(object_models, "_plane_tv_loss_eager", eager)
+    actual = object_models._plane_tv_loss(_pretend_cuda_grids(), tilted=True, rotations=2)
+    assert actual is expected
+    assert len(eager_calls) == 1
+    assert calls == []
diff --git a/tests/tomography/test_tomography.py b/tests/tomography/test_tomography.py
index da71f9c9e..b5ec106fa 100644
--- a/tests/tomography/test_tomography.py
+++ b/tests/tomography/test_tomography.py
@@ -12,12 +12,18 @@
 import pytest
 import torch
 
+from quantem.core.ml.optimizer_mixin import OptimizerParams, SchedulerParams
 from quantem.tomography.dataset_models import TomographyINRDataset, TomographyPixDataset
-from quantem.tomography.object_models import ObjConstraintParams, ObjectINR, ObjectPixelated
+from quantem.tomography.object_models import (
+    ObjConstraintParams,
+    ObjectINR,
+    ObjectPixelated,
+    ObjectTensorDecomp,
+)
 from quantem.tomography.tomography import Tomography, TomographyConventional
 from quantem.tomography.tomography_lite import TomographyLiteINR
 
-from .conftest import requires_torch
+from .conftest import requires_gpu, requires_torch
 
 
 def _stack(nang=5, n=12, seed=0):
@@ -141,6 +147,491 @@ def test_save_volume_overwrite_guard(self, torch_device, tmp_path):
         with np.load(path) as data:
             assert "volume" in data
 
+    def test_reconstruct_pose_warmup_activation_and_reference_exclusion(self, monkeypatch):
+        tomo = self._inr_tomo("cpu", n=4)
+        pose_active_at_epoch_start = []
+        pose_scheduler_at_epoch_start = []
+        object_active_at_epoch_start = []
+        initial_pose = [
+            param.detach().clone()
+            for param in (tomo.dset._shifts_params, tomo.dset._z1_params, tomo.dset._z3_params)
+        ]
+        pose_at_epoch_end = []
+        original_train = tomo.dset.train
+
+        def record_optimizer_state(mode=True):
+            pose_active_at_epoch_start.append(tomo.dset.has_optimizer())
+            pose_scheduler_at_epoch_start.append(tomo.dset.scheduler is not None)
+            object_active_at_epoch_start.append(tomo.obj_model.has_optimizer())
+            return original_train(mode)
+
+        def record_pose(_epoch):
+            pose_at_epoch_end.append(
+                [
+                    param.detach().clone()
+                    for param in (
+                        tomo.dset._shifts_params,
+                        tomo.dset._z1_params,
+                        tomo.dset._z3_params,
+                    )
+                ]
+            )
+
+        monkeypatch.setattr(tomo.dset, "train", record_optimizer_state)
+        tomo.reconstruct(
+            num_iter=2,
+            batch_size=len(tomo.dset),
+            num_workers=0,
+            num_samples_per_ray=2,
+            optimizer_params={
+                "object": OptimizerParams.Adam(lr=1e-3),
+                "pose": {
+                    "pose_shift": OptimizerParams.Adam(lr=2e-2),
+                    "pose_tilt_axis": OptimizerParams.Adam(lr=3e-3),
+                },
+            },
+            scheduler_params={
+                "object": SchedulerParams.Exponential(gamma=0.9),
+                "pose": SchedulerParams.Exponential(gamma=0.9),
+            },
+            pose_warmup_epochs=1,
+            eval_callback=record_pose,
+            eval_every=1,
+        )
+
+        assert pose_active_at_epoch_start == [False, True]
+        assert pose_scheduler_at_epoch_start == [False, True]
+        assert object_active_at_epoch_start == [True, True]
+        assert all(
+            torch.equal(before, after) for before, after in zip(initial_pose, pose_at_epoch_end[0])
+        )
+        assert any(
+            not torch.equal(before, after)
+            for before, after in zip(pose_at_epoch_end[0], pose_at_epoch_end[1])
+        )
+        assert {group["name"] for group in tomo.dset.optimizer.param_groups} == {
+            "pose_shift",
+            "pose_tilt_axis",
+        }
+        optimized = [
+            param for group in tomo.dset.optimizer.param_groups for param in group["params"]
+        ]
+        assert all(
+            ref is not param
+            for ref in (tomo.dset._shifts_ref, tomo.dset._z1_ref, tomo.dset._z3_ref)
+            for param in optimized
+        )
+
+        default_tomo = self._inr_tomo("cpu", n=4)
+        default_tomo.reconstruct(
+            num_iter=0,
+            batch_size=len(default_tomo.dset),
+            num_workers=0,
+            optimizer_params={"pose": OptimizerParams.Adam(lr=1e-2)},
+        )
+        assert default_tomo.dset.has_optimizer()
+
+        with pytest.raises(ValueError, match="pose_warmup_epochs must be >= 0"):
+            default_tomo.reconstruct(num_iter=0, pose_warmup_epochs=-1)
+
+    def test_reconstruct_float32_dtype_matches_disabled_autocast(self):
+        def run(autocast_dtype):
+            torch.manual_seed(17)
+            tomo = self._inr_tomo("cpu", n=4)
+            tomo.reconstruct(
+                num_iter=1,
+                batch_size=len(tomo.dset),
+                num_workers=0,
+                num_samples_per_ray=2,
+                autocast_dtype=autocast_dtype,
+            )
+            return tomo
+
+        disabled = run(None)
+        float32 = run(torch.float32)
+
+        np.testing.assert_allclose(float32.epoch_losses, disabled.epoch_losses)
+        for actual, expected in zip(
+            float32.obj_model.model.parameters(), disabled.obj_model.model.parameters()
+        ):
+            assert actual.dtype == torch.float32
+            torch.testing.assert_close(actual, expected)
+
+    @pytest.mark.parametrize("disabled_value", [None, 0.0])
+    def test_reconstruct_can_skip_inactive_gradient_clipping(self, monkeypatch, disabled_value):
+        original_clip = torch.nn.utils.clip_grad_norm_
+        clip_calls = []
+
+        def tracking_clip(parameters, max_norm, *args, **kwargs):
+            result = original_clip(parameters, max_norm, *args, **kwargs)
+            clip_calls.append((float(result), float(max_norm)))
+            return result
+
+        monkeypatch.setattr(torch.nn.utils, "clip_grad_norm_", tracking_clip)
+
+        def run(grad_clip_max_norm):
+            torch.manual_seed(19)
+            tomo = self._inr_tomo("cpu", n=4)
+            torch.manual_seed(23)
+            tomo.reconstruct(
+                num_iter=1,
+                batch_size=len(tomo.dset),
+                num_workers=0,
+                num_samples_per_ray=2,
+                grad_clip_max_norm=grad_clip_max_norm,
+            )
+            return (
+                tomo,
+                [
+                    parameter.grad.detach().clone()
+                    for parameter in tomo.obj_model.model.parameters()
+                ],
+                [parameter.detach().clone() for parameter in tomo.obj_model.model.parameters()],
+            )
+
+        clipped, clipped_grads, clipped_params = run(1e6)
+        unclipped, unclipped_grads, unclipped_params = run(disabled_value)
+
+        assert len(clip_calls) == 1
+        total_norm, max_norm = clip_calls[0]
+        assert total_norm < max_norm
+        np.testing.assert_allclose(unclipped.epoch_losses, clipped.epoch_losses)
+        for actual, expected in zip(unclipped_grads, clipped_grads):
+            torch.testing.assert_close(actual, expected)
+        for actual, expected in zip(unclipped_params, clipped_params):
+            torch.testing.assert_close(actual, expected)
+
+    def test_reconstruct_rejects_negative_gradient_clip_norm(self):
+        tomo = self._inr_tomo("cpu", n=4)
+        with pytest.raises(ValueError, match="grad_clip_max_norm"):
+            tomo.reconstruct(num_iter=0, grad_clip_max_norm=-1.0)
+
+    def test_reconstruct_rejects_invalid_autocast_dtype(self):
+        tomo = self._inr_tomo("cpu", n=4)
+        for invalid_dtype in ("float32", torch.float64):
+            with pytest.raises(ValueError, match="autocast_dtype"):
+                tomo.reconstruct(autocast_dtype=invalid_dtype)
+
+    def test_reconstruct_bfloat16_dtype_matches_bf16_string_on_cpu(self):
+        from quantem.core.ml.models.kplanes import KPlanes
+
+        def run(autocast_dtype):
+            torch.manual_seed(17)
+            n = 4
+            model = KPlanes(M_features=2, resolution=(n, n, n))
+            obj = ObjectTensorDecomp.from_model(model, shape=(n, n, n), device="cpu")
+            dset = TomographyINRDataset.from_data(
+                _stack(nang=5, n=n), np.linspace(-60, 60, 5).astype(np.float32)
+            )
+            tomo = Tomography.from_models(dset=dset, obj_model=obj, device="cpu", verbose=False)
+            tomo.reconstruct(
+                num_iter=2,
+                batch_size=len(dset),
+                num_workers=0,
+                num_samples_per_ray=2,
+                autocast_dtype=autocast_dtype,
+            )
+            return tomo
+
+        string_dtype = run("bf16")
+        torch_dtype = run(torch.bfloat16)
+
+        np.testing.assert_allclose(torch_dtype.epoch_losses, string_dtype.epoch_losses)
+        for actual, expected in zip(
+            torch_dtype.obj_model.model.parameters(),
+            string_dtype.obj_model.model.parameters(),
+        ):
+            torch.testing.assert_close(actual, expected)
+
+    @requires_gpu
+    def test_reconstruct_fp16_autocast_runs_on_cuda(self, monkeypatch):
+        grad_scalers = []
+        grad_scaler_cls = torch.amp.GradScaler
+
+        def capture_grad_scaler(*args, **kwargs):
+            grad_scaler = grad_scaler_cls(*args, **kwargs)
+            grad_scalers.append(grad_scaler)
+            return grad_scaler
+
+        monkeypatch.setattr(torch.amp, "GradScaler", capture_grad_scaler)
+        tomo = self._inr_tomo("cuda:0", n=4)
+        tomo.reconstruct(
+            num_iter=2,
+            batch_size=len(tomo.dset),
+            num_workers=0,
+            num_samples_per_ray=2,
+            autocast_dtype="fp16",
+        )
+
+        assert len(tomo.epoch_losses) == 2
+        assert np.isfinite(tomo.epoch_losses).all()
+        assert len(grad_scalers) == 1
+        scale = grad_scalers[0].get_scale()
+        assert isinstance(scale, float)
+        assert np.isfinite(scale) and scale > 0
+
+    @requires_gpu
+    @pytest.mark.parametrize(
+        ("model_kind", "multiscale_res_multipliers"),
+        [
+            pytest.param("inr", None, id="inr"),
+            pytest.param("kplanes_tilted", [1], id="kplanes-single-level"),
+            pytest.param(
+                "kplanes_tilted",
+                [0.25, 0.75, 1.0],
+                id="kplanes-three-level-ms-tv",
+            ),
+        ],
+    )
+    def test_reconstruct_cuda_graphs_matches_eager_fp32(
+        self, model_kind, multiscale_res_multipliers, monkeypatch
+    ):
+        seed = 7
+        monkeypatch.delenv("QUANTEM_KPLANES_MS_TV_FUSED", raising=False)
+
+        def build_tomography():
+            if model_kind == "inr":
+                return self._inr_tomo("cuda:0", n=4)
+
+            from quantem.core.ml.models.kplanes import KPlanesTILTED
+
+            model = KPlanesTILTED(
+                M_features=2,
+                resolution=(4, 4, 4),
+                multiscale_res_multipliers=multiscale_res_multipliers,
+                T=2,
+                so3_param_type="r9svd",
+            )
+            obj = ObjectTensorDecomp.from_model(
+                model,
+                shape=(4, 4, 4),
+                device="cuda:0",
+            )
+            if len(multiscale_res_multipliers) == 3:
+                obj.constraints.tv_plane = 0.05
+            dset = TomographyINRDataset.from_data(
+                _stack(nang=5, n=4),
+                np.linspace(-60, 60, 5).astype(np.float32),
+            )
+            return Tomography.from_models(
+                dset=dset,
+                obj_model=obj,
+                device="cuda:0",
+                verbose=False,
+            )
+
+        torch.manual_seed(seed)
+        eager = build_tomography()
+        torch.manual_seed(seed)
+        graphed = build_tomography()
+
+        eager_initial_params = [
+            parameter.detach().clone() for parameter in eager.obj_model.model.parameters()
+        ]
+        graphed_initial_params = [
+            parameter.detach().clone() for parameter in graphed.obj_model.model.parameters()
+        ]
+
+        reconstruct_kwargs = {
+            "num_iter": 40,
+            "batch_size": len(eager.dset),
+            "num_workers": 0,
+            "num_samples_per_ray": 4,
+        }
+        if model_kind == "inr":
+            reconstruct_kwargs["optimizer_params"] = {"object": OptimizerParams.Adam(lr=1e-3)}
+        else:
+            optimizer_params = {
+                "grids": OptimizerParams.Adam(lr=1e-3),
+                "sigma_net": OptimizerParams.Adam(lr=1e-3),
+                "so3": OptimizerParams.Adam(lr=1e-3),
+            }
+            eager.obj_model.set_optimizer(optimizer_params)
+            graphed.obj_model._cuda_graphs_optimizer = True
+            try:
+                graphed.obj_model.set_optimizer(optimizer_params)
+            finally:
+                graphed.obj_model._cuda_graphs_optimizer = False
+            eager.obj_model.model.so3.requires_grad_(False)
+            graphed.obj_model.model.so3.requires_grad_(False)
+
+        def adam_steps(tomo):
+            return [
+                int(state["step"].item())
+                for state in tomo.obj_model.optimizer.state.values()
+                if "step" in state
+            ]
+
+        graphed_step_history = []
+
+        def record_graphed_step():
+            steps = adam_steps(graphed)
+            assert steps and len(set(steps)) == 1
+            graphed_step_history.append(steps[0])
+
+        graphed._cuda_graph_step_callback = record_graphed_step
+
+        torch.manual_seed(seed)
+        eager.reconstruct(**reconstruct_kwargs)
+        torch.manual_seed(seed)
+        graphed.reconstruct(**reconstruct_kwargs, cuda_graphs=True)
+        if model_kind == "kplanes_tilted":
+            assert graphed.obj_model.model._rotation_matrices_override is None
+
+        eager_param_delta = max(
+            (parameter - initial).abs().max().item()
+            for parameter, initial in zip(eager.obj_model.model.parameters(), eager_initial_params)
+        )
+        graphed_param_delta = max(
+            (parameter - initial).abs().max().item()
+            for parameter, initial in zip(
+                graphed.obj_model.model.parameters(), graphed_initial_params
+            )
+        )
+        assert eager_param_delta > 0.0
+        assert graphed_param_delta > 0.0
+        assert graphed_step_history == list(range(1, reconstruct_kwargs["num_iter"] + 1))
+
+        eager_steps = adam_steps(eager)
+        graphed_steps = adam_steps(graphed)
+        assert eager_steps
+        assert graphed_steps
+        assert set(eager_steps) == {reconstruct_kwargs["num_iter"]}
+        assert graphed_steps == eager_steps
+
+        # Padded rays change reduction order, so Adam compounds small graph/eager
+        # differences without a bound in principle. The full trajectory is only a smoke
+        # bound; the prefix check and Adam-step/parameter-delta assertions establish
+        # faithfulness.
+        np.testing.assert_allclose(
+            graphed.epoch_losses[:10],
+            eager.epoch_losses[:10],
+            rtol=1e-6,
+            atol=0.0,
+        )
+        np.testing.assert_allclose(
+            graphed.epoch_losses,
+            eager.epoch_losses,
+            rtol=2e-3,
+            atol=0.0,
+        )
+
+    @requires_gpu
+    @pytest.mark.parametrize("num_steps", [1, 10])
+    def test_pred_fork_matches_single_stream_state(self, monkeypatch, num_steps):
+        def build_tomography():
+            tomo = self._inr_tomo("cuda:0", n=4)
+            tomo.obj_model.constraints = ObjConstraintParams.ObjINRConstraints(
+                s3im_weight=0.05,
+                s3im_repeat_time=2,
+                s3im_kernel=2,
+                s3im_value_range=10.0,
+            )
+            return tomo
+
+        def run(tomo, enabled):
+            if enabled:
+                monkeypatch.setenv("QUANTEM_RECON_PRED_FORK", "1")
+            else:
+                monkeypatch.setenv("QUANTEM_RECON_PRED_FORK", "0")
+            torch.manual_seed(29)
+            tomo.reconstruct(
+                num_iter=num_steps,
+                batch_size=len(tomo.dset),
+                num_workers=0,
+                num_samples_per_ray=2,
+                optimizer_params={"object": OptimizerParams.Adam(lr=1e-3)},
+                grad_clip_max_norm=None,
+            )
+            torch.cuda.synchronize()
+
+        torch.manual_seed(17)
+        reference = build_tomography()
+        torch.manual_seed(17)
+        forked = build_tomography()
+        run(reference, enabled=False)
+        run(forked, enabled=True)
+
+        # The 100-step arm was removed because intrinsic atomic-order chaos
+        # (1.5e-2 ref-vs-ref) exceeds any assertable bound; the standing measurement is
+        # test_pred_fork_reference_intrinsic_drift.
+        np.testing.assert_allclose(forked.epoch_losses, reference.epoch_losses, rtol=1e-4)
+        np.testing.assert_allclose(
+            forked.consistency_losses, reference.consistency_losses, rtol=1e-4
+        )
+        np.testing.assert_allclose(
+            forked.obj_model.soft_constraint_losses,
+            reference.obj_model.soft_constraint_losses,
+            rtol=1e-4,
+        )
+        for forked_parameter, reference_parameter in zip(
+            forked.obj_model.model.parameters(), reference.obj_model.model.parameters()
+        ):
+            torch.testing.assert_close(
+                forked_parameter.grad,
+                reference_parameter.grad,
+                rtol=1e-4,
+                atol=1e-7,
+            )
+            torch.testing.assert_close(
+                forked_parameter,
+                reference_parameter,
+                rtol=1e-4,
+                atol=1e-7,
+            )
+            forked_state = forked.obj_model.optimizer.state[forked_parameter]
+            reference_state = reference.obj_model.optimizer.state[reference_parameter]
+            assert forked_state.keys() == reference_state.keys()
+            for key in forked_state:
+                if isinstance(forked_state[key], torch.Tensor):
+                    torch.testing.assert_close(
+                        forked_state[key],
+                        reference_state[key],
+                        rtol=1e-4,
+                        atol=1e-7,
+                    )
+                else:
+                    assert forked_state[key] == reference_state[key]
+
+    @requires_gpu
+    @pytest.mark.parametrize("num_steps", [10, 100])
+    def test_pred_fork_reference_intrinsic_drift(self, monkeypatch, num_steps):
+        def build_tomography():
+            tomo = self._inr_tomo("cuda:0", n=4)
+            tomo.obj_model.constraints = ObjConstraintParams.ObjINRConstraints(
+                s3im_weight=0.05,
+                s3im_repeat_time=2,
+                s3im_kernel=2,
+                s3im_value_range=10.0,
+            )
+            return tomo
+
+        def run(tomo):
+            monkeypatch.setenv("QUANTEM_RECON_PRED_FORK", "0")
+            torch.manual_seed(29)
+            tomo.reconstruct(
+                num_iter=num_steps,
+                batch_size=len(tomo.dset),
+                num_workers=0,
+                num_samples_per_ray=2,
+                optimizer_params={"object": OptimizerParams.Adam(lr=1e-3)},
+                grad_clip_max_norm=None,
+            )
+            torch.cuda.synchronize()
+
+        torch.manual_seed(17)
+        reference_a = build_tomography()
+        torch.manual_seed(17)
+        reference_b = build_tomography()
+        run(reference_a)
+        run(reference_b)
+
+        losses_a = reference_a.epoch_losses
+        losses_b = reference_b.epoch_losses
+        d = np.max(np.abs(losses_a - losses_b) / np.abs(losses_a))
+        print(f"intrinsic drift @{num_steps} steps: {d:.3e}")
+        assert np.isfinite(d)
+
 
 @requires_torch
 class TestLiteINRReconstructBranch:
diff --git a/tests/tomography/test_tomography_conventional.py b/tests/tomography/test_tomography_conventional.py
index a7cfb5e02..85e3daf23 100644
--- a/tests/tomography/test_tomography_conventional.py
+++ b/tests/tomography/test_tomography_conventional.py
@@ -61,6 +61,27 @@ def test_obj_constraints_accepts_dict(self, phantom_volume, tilt_series, tilt_an
         assert tomo.num_epochs == 2
 
 
+class TestInlineAlignment:
+    def test_alignment_corrects_misaligned_projection(
+        self, phantom_volume, tilt_series, tilt_angles
+    ):
+        """Regression: the aligned measurement was written into proj_forward, which
+        radon_torch overwrites immediately, so inline_alignment was a no-op."""
+        n = phantom_volume.shape[0]
+        shifted_series = tilt_series.copy()
+        shifted_series[4] = np.roll(shifted_series[4], shift=3, axis=1)
+
+        reference = _build(tilt_series, tilt_angles, n).dset.tilt_stack[4]
+        tomo = _build(shifted_series, tilt_angles, n)
+        misaligned_before = tomo.dset.tilt_stack[4].clone()
+        tomo.reconstruct(num_iter=3, mode="sirt", inline_alignment=True)
+        aligned_after = tomo.dset.tilt_stack[4]
+
+        err_before = float(((misaligned_before - reference) ** 2).mean())
+        err_after = float(((aligned_after - reference) ** 2).mean())
+        assert err_after < err_before  # the stack was actually re-aligned
+
+
 class TestFBP:
     def test_fbp_runs_single_epoch(self, phantom_volume, tilt_series, tilt_angles):
         n = phantom_volume.shape[0]
diff --git a/tests/tomography/test_utils.py b/tests/tomography/test_utils.py
index 1fa6a47ce..b85304912 100644
--- a/tests/tomography/test_utils.py
+++ b/tests/tomography/test_utils.py
@@ -5,6 +5,7 @@
 import torch
 
 from quantem.tomography.utils import (
+    differentiable_rotx_vectorized,
     differentiable_rotz_vectorized,
     rot_ZXZ,
     tv_loss_1d,
@@ -76,3 +77,37 @@ def test_gradient_flows_through_rotation(self):
         out.sum().backward()
         assert vol.grad is not None
         assert torch.isfinite(vol.grad).all()
+
+
+class TestRotZXZGradients:
+    def test_grad_flows_with_mixed_float_and_tensor_angles(self):
+        """Regression: a non-tensor angle made rot_ZXZ re-wrap every angle with
+        torch.tensor(), detaching gradients through tensor angles."""
+        vol = _block_volume()
+        x = torch.tensor(20.0, requires_grad=True)
+        out = rot_ZXZ(vol, 0.0, x, 0.0, device="cpu")
+        out.sum().backward()
+        assert x.grad is not None
+        assert torch.isfinite(x.grad)
+    @pytest.mark.parametrize(
+        "rot_fn", [differentiable_rotz_vectorized, differentiable_rotx_vectorized]
+    )
+    def test_multi_angle_matches_per_angle_calls(self, rot_fn):
+        # Regression: the per-slice vmap version raised for more than one angle
+        # (affine_grid requires the matrix batch to match the slice batch of 1).
+        vol = _block_volume()
+        angles = torch.tensor([10.0, -25.0, 60.0])
+        multi = rot_fn(vol, angles)
+        per_angle = torch.cat([rot_fn(vol, a) for a in angles])
+        assert multi.shape == (3, *vol.shape[1:])
+        assert torch.allclose(multi, per_angle, atol=1e-6)
+
+    @pytest.mark.parametrize(
+        "rot_fn", [differentiable_rotz_vectorized, differentiable_rotx_vectorized]
+    )
+    def test_per_volume_angles(self, rot_fn):
+        vols = torch.cat([_block_volume(), _block_volume().flip(-1), _block_volume().flip(-2)])
+        angles = torch.tensor([10.0, -25.0, 60.0])
+        batched = rot_fn(vols, angles)
+        per_volume = torch.cat([rot_fn(vols[i : i + 1], angles[i]) for i in range(3)])
+        assert torch.allclose(batched, per_volume, atol=1e-6)
diff --git a/tests/tomography/test_validation_loss.py b/tests/tomography/test_validation_loss.py
new file mode 100644
index 000000000..e5d9fbcf9
--- /dev/null
+++ b/tests/tomography/test_validation_loss.py
@@ -0,0 +1,104 @@
+import torch
+from torch import nn
+
+from quantem.tomography.tomography import Tomography
+
+
+class _CountingValidationLoader:
+    def __init__(self):
+        self.iter_calls = 0
+        self.batch_counts: list[int] = []
+
+    def __iter__(self):
+        self.iter_calls += 1
+        batches = [
+            {
+                "x": torch.tensor([[1.0], [2.0]]),
+                "target_value": torch.tensor([0.0, 0.0]),
+            },
+            {
+                "x": torch.tensor([[3.0]]),
+                "target_value": torch.tensor([0.0]),
+            },
+        ]
+        self.batch_counts.append(len(batches))
+        return iter(batches)
+
+
+class _ValidationDataset(nn.Module):
+    def get_coords(
+        self, batch: dict[str, torch.Tensor], object_extent: int, num_samples_per_ray: int
+    ) -> torch.Tensor:
+        return batch["x"]
+
+    def integrate_rays(
+        self, densities: torch.Tensor, num_samples_per_ray: int, target_values_len: int
+    ) -> torch.Tensor:
+        return densities.reshape(target_values_len)
+
+
+class _ValidationObject(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.model = nn.Linear(1, 1, bias=False)
+
+    def forward(self, coords: torch.Tensor) -> torch.Tensor:
+        return self.model(coords).squeeze(-1)
+
+
+def _tomography_for_validation() -> Tomography:
+    tomography = object.__new__(Tomography)
+    tomography.device = torch.device("cpu")
+    tomography.world_size = 1
+    tomography._dset = _ValidationDataset()
+    tomography._obj_model = _ValidationObject()
+    return tomography
+
+
+def _evaluate(tomography: Tomography, loader: _CountingValidationLoader) -> float:
+    loss = tomography._evaluate_validation_loss(
+        dataloader=loader,
+        num_samples_per_ray=1,
+        object_extent=1,
+        loss_func=nn.MSELoss(),
+    )
+    assert loss is not None
+    return loss
+
+
+def test_validation_loss_depends_on_current_model_parameters():
+    tomography = _tomography_for_validation()
+    loader = _CountingValidationLoader()
+
+    with torch.no_grad():
+        tomography.obj_model.model.weight.fill_(1.0)
+    first_loss = _evaluate(tomography, loader)
+
+    with torch.no_grad():
+        tomography.obj_model.model.weight.fill_(2.0)
+    second_loss = _evaluate(tomography, loader)
+
+    assert first_loss != second_loss
+
+
+def test_validation_loss_reiterates_fresh_dataloader_each_call():
+    tomography = _tomography_for_validation()
+    loader = _CountingValidationLoader()
+
+    _evaluate(tomography, loader)
+    _evaluate(tomography, loader)
+
+    assert loader.iter_calls == 2
+    assert loader.batch_counts == [2, 2]
+
+
+def test_validation_loss_restores_train_mode():
+    tomography = _tomography_for_validation()
+    loader = _CountingValidationLoader()
+    tomography.obj_model.model.train()
+    tomography.dset.train()
+
+    _evaluate(tomography, loader)
+
+    assert tomography.obj_model.model.training
+    assert tomography.dset.training