diffct 2.4.4__py3-none-any.whl → 2.5.0__py3-none-any.whl

diffct/__init__.py

@@ -64,7 +64,7 @@ try:
 except ImportError:
     pass
 
-__version__ = '2.4.4'
+__version__ = '2.5.0'
 
 __all__ = [
     'Projector',

diffct/_solver_state.py

@@ -12,6 +12,7 @@ import torch.distributed as dist
 from .pipeline import _fragments
 from .storage import NpyStore, TensorStore
 from .distribution import control_status, preflight
+from .utils import _work_dtype
 
 
 def count(value, name, positive=False):
@@ -65,20 +66,46 @@ def scalar_sum(operator, value, kind, *, return_tensor=False):
         dist.all_reduce(accumulator, group=group)
         return accumulator if return_tensor else accumulator.item()
     if return_tensor:
-        return torch.as_tensor(value, dtype=torch.float64, device=result_device)
+        return torch.as_tensor(value, dtype=_work_dtype(result_device), device=result_device)
     return float(value)
 
 
 def squared_norm(operator, value, kind, elements=1048576, *, return_tensor=False):
     store = TensorStore(value) if isinstance(value, torch.Tensor) else value
-    total = torch.zeros((), dtype=torch.float64, device=getattr(store, "device", "cpu"))
+    # Scale every block by its own peak before squaring and track the
+    # running maximum on the device: unscaled FP32 squares overflow for
+    # large measurements and underflow for tiny geometries.
+    device = getattr(store, "device", torch.device("cpu"))
+    if _work_dtype(device) == torch.float64:
+        total = torch.zeros((), dtype=torch.float64, device=device)
+        if owns_values(operator, kind):
+            for selected in _fragments(tuple(store.shape), elements):
+                # Clone even an FP64 view: norms must never alter caller measurements.
+                block = store.read(selected).to(dtype=torch.float64).clone()
+                # Preserve the block reductions and their sequential FP64 order.
+                total.add_(block.square_().sum().to(device=total.device))
+        return scalar_sum(operator, total, kind, return_tensor=return_tensor)
+    # FP32 devices (MPS): squares of large measurements overflow and tiny
+    # geometries underflow, so every block is rescaled by its peak and the
+    # running sum follows the largest peak seen. All block math stays on the
+    # device; the two closing scalar reads fit the per-iteration reduction
+    # budget, and their product leaves FP32 before it can overflow.
+    peak = torch.zeros((), device=device, dtype=torch.float32)
+    total = torch.zeros((), device=device, dtype=torch.float32)
     if owns_values(operator, kind):
         for selected in _fragments(tuple(store.shape), elements):
-            # Clone even an FP64 view: norms must never alter caller measurements.
-            block = store.read(selected).to(dtype=torch.float64).clone()
-            # Preserve the block reductions and their sequential FP64 order.
-            total.add_(block.square_().sum().to(device=total.device))
-    return scalar_sum(operator, total, kind, return_tensor=return_tensor)
+            block = store.read(selected)
+            if not block.numel():
+                continue
+            block_peak = block.detach().abs().amax()
+            safe = torch.where(block_peak > 0, block_peak, torch.ones_like(block_peak))
+            combined = torch.maximum(peak, block_peak)
+            guard = torch.where(combined > 0, combined, torch.ones_like(combined))
+            total = (total * (peak / guard).square()
+                     + (block.detach() / safe).square().sum() * (block_peak / guard).square())
+            peak = combined
+    return scalar_sum(operator, peak.item() ** 2 * total.item(), kind,
+                      return_tensor=return_tensor)
 
 
 def copy_blocks(source, target, elements):

diffct/analytical_execution.py

@@ -160,6 +160,8 @@ class AnalyticalExecution:
             raise TypeError("sinogram store must have dtype float32 or float64")
         self.source = source
         self.output_device = torch.device(getattr(source, "device", "cpu"))
+        if self.dtype == torch.float64 and self.output_device.type == "mps":
+            raise ValueError("precision must be 'fast' or 'float32' on MPS; MPS provides no float64")
         self.detector_shape = tuple(source.shape[1:])
         for component in trajectory:
             _fixed(component, "trajectory")

diffct/backends.py

@@ -12,6 +12,7 @@ import torch.distributed as dist
 from torch.autograd import forward_ad
 
 from .chunking import _memory_lower_bounds
+from .utils import _work_dtype
 from .distribution import control_status
 from .surfaces import ParameterizedSurface
 from .constants import _native_spacing
@@ -76,6 +77,16 @@ def _zero_dependency(value):
     return value.unsqueeze(0)[:0].sum()
 
 
+def _constant_like(reference, values):
+    """Build a constant tensor on ``reference``'s device and dtype.
+
+    ``new_tensor`` has no kernel behind a functorch dual on MPS and raises
+    "DispatchKey Undefined doesn't correspond to a device"; an explicit
+    constructor keeps the same constant on every device and transform.
+    """
+    return torch.tensor(values, dtype=reference.dtype, device=reference.device)
+
+
 def voxel_spacing_tuple(spacing, rank):
     """Return positive physical pitches in tensor axis order."""
     values = spacing if isinstance(spacing, (tuple, list)) else (spacing,) * rank
@@ -563,7 +574,14 @@ class ReferenceProjector:
             with torch.no_grad():
                 # Reuse frame transfers across detector chunks. Offset batches
                 # remain bounded and move only when a ray batch needs them.
-                staged = {device: (*tuple(value.to(device=device) for value in geometry[:frame_count]),
+                def _stage(value, target_device):
+                    # FP64 frames must not follow FP64-less compute devices.
+                    target_dtype = _work_dtype(target_device)
+                    if value.dtype.itemsize > target_dtype.itemsize:
+                        return value.to(device=target_device, dtype=target_dtype)
+                    return value.to(device=target_device)
+
+                staged = {device: (*tuple(_stage(value, device) for value in geometry[:frame_count]),
                                     *parameters) for device in set(devices)}
                 batches = (self._admission_batches(staged[devices[0]], plan) if admission
                            else self._batches(*plan[1:]))
@@ -630,7 +648,7 @@ class ReferenceProjector:
         def compose(validate, stats):
             check_range = (validate and owner.beam not in ("parallel", "parallel3d")
                            and owner.dtype == torch.float32 and surface is not None)
-            dtype = torch.float64 if check_range else owner.dtype
+            dtype = _work_dtype(device) if check_range else owner.dtype
             frames = (owner._learnable or owner._trajectory) if geometry is None else geometry[:frame_count]
             native_surface = owner.backend == "numba" and surface is not None
             native_legacy = native_surface and not isinstance(surface, ParameterizedSurface)
@@ -666,7 +684,7 @@ class ReferenceProjector:
             if check_range and not native_surface:
                 # Certify ordinary batches without expanding raw FP64 world points.
                 # L1 norms bound Euclidean distances without squaring tiny bases.
-                pitches = local.new_tensor(self.spacing[::-1])
+                pitches = _constant_like(local, self.spacing[::-1])
                 maximum = local.abs().amax(dim=tuple(range(1, local.ndim - 1)))
                 axis_u = selected[2]
                 if owner.beam == "cone":
@@ -698,7 +716,7 @@ class ReferenceProjector:
             if check_range:
                 # Raw FP64 coordinates preserve the exact admission boundaries.
                 sources = frame(selected[0]).expand_as(points)
-                pitches = points.new_tensor(self.spacing[::-1])
+                pitches = _constant_like(points, self.spacing[::-1])
                 source_distance = torch.linalg.vector_norm(sources / pitches, dim=-1)
                 pixel_distance = torch.linalg.vector_norm(points / pitches, dim=-1)
                 if torch.any(torch.minimum(source_distance, pixel_distance) > 1e6).item():
@@ -725,7 +743,7 @@ class ReferenceProjector:
                 origins, directions = points, frame(selected[0]).expand_as(points)
             else:
                 sources = frame(selected[0]).expand_as(points)
-                pitches = points.new_tensor(self.spacing[::-1])
+                pitches = _constant_like(points, self.spacing[::-1])
                 source_distance = torch.linalg.vector_norm(sources / pitches, dim=-1)
                 pixel_distance = torch.linalg.vector_norm(points / pitches, dim=-1)
                 if validate:
@@ -764,14 +782,14 @@ class ReferenceProjector:
         # Global FP32 segment parameters lose voxel widths when both endpoints
         # are distant. Plane tables below use local physical distances and the
         # chosen precision, including for native data-transform dispatch.
-        origins = origins.to(dtype=torch.float64)
-        directions = directions.to(dtype=torch.float64)
+        origins = origins.to(dtype=_work_dtype(origins.device))
+        directions = directions.to(dtype=_work_dtype(directions.device))
         ray_length = torch.linalg.vector_norm(directions, dim=-1)
         directions = directions / ray_length[:, None]
-        lower = origins.new_tensor([(part.start - size / 2) * pitch
-                                    for part, size, pitch in zip(world_tile, shape, spacing)])
-        upper = origins.new_tensor([(part.stop - size / 2) * pitch
-                                    for part, size, pitch in zip(world_tile, shape, spacing)])
+        lower = _constant_like(origins, [(part.start - size / 2) * pitch
+                                        for part, size, pitch in zip(world_tile, shape, spacing)])
+        upper = _constant_like(origins, [(part.stop - size / 2) * pitch
+                                        for part, size, pitch in zip(world_tile, shape, spacing)])
         constant = directions == 0
         safe_direction = torch.where(constant, torch.ones_like(directions), directions)
         first, last = (lower - origins) / safe_direction, (upper - origins) / safe_direction
@@ -1093,7 +1111,17 @@ class ReferenceProjector:
             # The zero dependency keeps unused inputs connected. With several
             # devices, a pair's loss and the geometry leaves can be on different
             # CUDA devices, so each zero moves to the loss's device.
-            loss = loss + sum(_zero_dependency(value).to(loss.device) for value in inputs)
+            def _zero_to(value, device):
+                zero = _zero_dependency(value)
+                target_dtype = _work_dtype(device)
+                if zero.dtype.itemsize > target_dtype.itemsize:
+                    # Cast on the source device first: a fused device and dtype
+                    # copy routes the float64 gradient back through the
+                    # FP64-less device on the way out.
+                    return zero.to(dtype=target_dtype).to(device=device)
+                return zero.to(device=device)
+
+            loss = loss + sum(_zero_to(value, loss.device) for value in inputs)
             for total, part in zip(totals, torch.autograd.grad(loss, inputs)):
                 total.add_(part)
 

diffct/geometry.py

@@ -5,6 +5,7 @@ and custom trajectories for parallel, fan and cone beams, plus rigid transforms.
 """
 
 import math
+import operator
 import torch
 from typing import Callable, Tuple
 
@@ -40,10 +41,20 @@ def transform_trajectory(trajectory, rotation=None, translation=None, *, beam='c
     if rotation is not None:
         if tuple(rotation.shape) not in ((dimensions, dimensions), (views, dimensions, dimensions)):
             raise ValueError("rotation must have shape (d, d) or (views, d, d)")
-        matrix = rotation.detach().double()
+        # MPS provides no float64; the formal harness also loads this module
+        # standalone, so the cast stays inline instead of importing helpers.
+        matrix = rotation.detach().to(dtype=torch.float32 if rotation.device.type == "mps"
+                                      else torch.float64)
         identity = torch.eye(dimensions, device=matrix.device, dtype=matrix.dtype).expand_as(matrix)
+        # Small explicit determinants: linalg.det routes through LU kernels
+        # that MPS does not provide.
+        if dimensions == 2:
+            determinant = (matrix[..., 0, 0] * matrix[..., 1, 1]
+                           - matrix[..., 0, 1] * matrix[..., 1, 0])
+        else:
+            determinant = (matrix[..., 0] * torch.linalg.cross(matrix[..., 1], matrix[..., 2], dim=-1)).sum(-1)
         if (not torch.allclose(matrix.transpose(-1, -2) @ matrix, identity, rtol=1e-4, atol=1e-4)
-                or not torch.allclose(torch.linalg.det(matrix), matrix.new_ones(matrix.shape[:-2]),
+                or not torch.allclose(determinant, matrix.new_ones(matrix.shape[:-2]),
                                       rtol=1e-4, atol=1e-4)):
             raise ValueError("rotation must be orthogonal and proper (determinant +1)")
     if translation is not None and tuple(translation.shape) not in ((dimensions,), (views, dimensions)):
@@ -84,22 +95,50 @@ def _circular_angles(n_views, start_angle, end_angle, device, dtype):
         if not torch.isfinite(n_views).all().item():
             raise ValueError("angles must be finite")
         return n_views, n_views.numel(), n_views.device, n_views.dtype
+    _check_n_views(n_views)
     step = (end_angle - start_angle) / n_views
     angles = start_angle + torch.arange(n_views, device=device, dtype=dtype) * step
     return angles, n_views, device, dtype
 
 
-def _check_source_positions(src_pos):
+def _check_source_positions(src_pos, what="source_path_fn"):
     """Reject non-finite sources and sources at the isocenter, which have no viewing direction."""
     if not torch.isfinite(src_pos).all():
-        raise ValueError("source_path_fn returned non-finite source positions")
+        raise ValueError(f"{what} returned non-finite source positions")
     if (torch.linalg.vector_norm(src_pos.detach(), dim=1) == 0).any():
         raise ValueError(
-            "source_path_fn returned a source at the isocenter; every source needs a "
+            f"{what} returned a source at the isocenter; every source needs a "
             "nonzero distance from the origin"
         )
 
 
+def _check_n_views(n_views):
+    """Reject view counts that are not positive integers.
+
+    A float count would space ``arange(ceil(n))`` views by ``span / n``, and a
+    nonpositive count builds an empty trajectory or divides by zero.
+    """
+    try:
+        count = operator.index(n_views)
+    except TypeError:
+        count = None
+    if count is None or isinstance(n_views, bool):
+        raise TypeError(f"n_views must be a positive integer, got {n_views!r}")
+    if count < 1:
+        raise ValueError(f"n_views must be a positive integer, got {n_views!r}")
+
+
+def _check_ray_directions(ray_dir):
+    """Reject non-finite and zero ray directions, which normalize to NaN."""
+    if not torch.isfinite(ray_dir).all():
+        raise ValueError("ray_dir_fn returned non-finite ray directions")
+    if (torch.linalg.vector_norm(ray_dir.detach(), dim=1) == 0).any():
+        raise ValueError(
+            "ray_dir_fn returned a zero ray direction; every view needs a "
+            "nonzero direction vector"
+        )
+
+
 def circular_trajectory_3d(n_views, sid, sdd, start_angle=0.0, end_angle=None, device='cuda', dtype=torch.float32):
     """Generate circular trajectory geometry for cone-beam CT.
 
@@ -286,6 +325,8 @@ def random_trajectory_3d(n_views, sid_mean, sdd_mean, sid_std=0.0, pos_std=0.0,
     det_u_vec = det_u_vec / torch.norm(det_u_vec, dim=1, keepdim=True)
     det_v_vec = det_v_vec / torch.norm(det_v_vec, dim=1, keepdim=True)
 
+    _check_source_positions(src_pos, "random_trajectory_3d")
+
     return src_pos, det_center, det_u_vec, det_v_vec
 
 
@@ -334,6 +375,7 @@ def spiral_trajectory_3d(n_views, sid, sdd, z_range=100.0, n_turns=2.0,
     """
     # Generate angles for spiral motion
     end_angle = start_angle + 2 * math.pi * n_turns
+    _check_n_views(n_views)
     step = (end_angle - start_angle) / n_views
     angles = start_angle + torch.arange(n_views, device=device, dtype=dtype) * step
 
@@ -419,6 +461,7 @@ def sinusoidal_trajectory_3d(n_views, sid, sdd, amplitude=50.0, frequency=2.0,
     """
     # Generate angles
     end_angle = start_angle + 2 * math.pi
+    _check_n_views(n_views)
     step = (end_angle - start_angle) / n_views
     angles = start_angle + torch.arange(n_views, device=device, dtype=dtype) * step
 
@@ -505,6 +548,7 @@ def saddle_trajectory_3d(n_views, sid, sdd, z_amplitude=50.0, radial_amplitude=3
     """
     # Generate angles
     end_angle = start_angle + 2 * math.pi
+    _check_n_views(n_views)
     step = (end_angle - start_angle) / n_views
     angles = start_angle + torch.arange(n_views, device=device, dtype=dtype) * step
 
@@ -601,6 +645,7 @@ def custom_trajectory_3d(n_views, sid, sdd,
     """
     # Generate angles
     end_angle = start_angle + 2 * math.pi
+    _check_n_views(n_views)
     step = (end_angle - start_angle) / n_views
     angles = start_angle + torch.arange(n_views, device=device, dtype=dtype) * step
 
@@ -746,6 +791,7 @@ def sinusoidal_trajectory_2d_fan(n_views, sid, sdd, amplitude=50.0, frequency=2.
     """
     # Generate angles
     end_angle = start_angle + 2 * math.pi
+    _check_n_views(n_views)
     step = (end_angle - start_angle) / n_views
     angles = start_angle + torch.arange(n_views, device=device, dtype=dtype) * step
 
@@ -829,6 +875,7 @@ def custom_trajectory_2d_fan(n_views, sid, sdd,
     """
     # Generate angles
     end_angle = start_angle + 2 * math.pi
+    _check_n_views(n_views)
     step = (end_angle - start_angle) / n_views
     angles = start_angle + torch.arange(n_views, device=device, dtype=dtype) * step
 
@@ -973,6 +1020,7 @@ def sinusoidal_trajectory_2d_parallel(n_views, amplitude=50.0, frequency=2.0,
     """
     # Generate angles
     end_angle = start_angle + math.pi
+    _check_n_views(n_views)
     step = (end_angle - start_angle) / n_views
     angles = start_angle + torch.arange(n_views, device=device, dtype=dtype) * step
 
@@ -1055,6 +1103,7 @@ def custom_trajectory_2d_parallel(n_views,
     """
     # Generate angles
     end_angle = start_angle + math.pi
+    _check_n_views(n_views)
     step = (end_angle - start_angle) / n_views
     angles = start_angle + torch.arange(n_views, device=device, dtype=dtype) * step
 
@@ -1069,6 +1118,9 @@ def custom_trajectory_2d_parallel(n_views,
         raise ValueError(f"det_origin_fn must return tensor of shape ({n_views}, 2), "
                         f"got {det_origin.shape}")
 
+    # Reject degenerate directions before normalization turns them into NaN
+    _check_ray_directions(ray_dir)
+
     # Normalize ray directions
     ray_dir = ray_dir / torch.norm(ray_dir, dim=1, keepdim=True)
 

diffct/operators.py

@@ -9,6 +9,7 @@ import warnings
 
 import torch
 
+from .constants import _native_spacing
 from .chunking import (_ChunkPlan, _allocation_bytes, _host_geometry_bytes, _largest_buffer_bytes, _memory_budget,
                        _memory_lower_bounds, _pipeline_buffers, _pipeline_bytes, _pixel_tiles,
                        _reservation_budget, _reservation_bytes, _tiles, _working_set_bytes)
@@ -17,11 +18,10 @@ from .storage import TensorStore
 from .streaming import _validate_stores
 from .surfaces import ParameterizedSurface, _GeometrySpec, _sample_geometry, _sample_offsets
 from .surface_gradients import surface_batch_vjp
-from .utils import TorchCUDABridge, _get_numba_external_stream_for, _raise_geometry_errors
+from .utils import TorchCUDABridge, _cast_work, _get_numba_external_stream_for, _raise_geometry_errors, _work_dtype
 from .backends import (_GEOMETRY_FORWARD_DERIVATIVE, ReferenceProjector, _flat_admission_batches, _forward_tangent,
                        select_backend, voxel_spacing_tuple)
 from .distribution import Partition
-from .constants import _native_spacing
 from .projectors import (
     _backproject_into,
     _geometry_vjp,
@@ -126,28 +126,36 @@ def _validate_trajectory(trajectory, beam, *, detector_surface=False, clone=True
         dots = torch.sum(first.to(device) * second.to(device), dim=1)
         checks.append((torch.any(torch.abs(dots) > 1e-4), "trajectory direction axes must be orthogonal"))
     if beam == "parallel3d":
-        direction = geometry[0].to(device=device, dtype=torch.float64)
-        frame = torch.stack((direction, geometry[2].to(direction.device, torch.float64),
-                             geometry[3].to(direction.device, torch.float64)), dim=-1)
+        direction = _cast_work(geometry[0], device)
+        frame = torch.stack((direction, _cast_work(geometry[2], direction.device),
+                             _cast_work(geometry[3], direction.device)), dim=-1)
         # Numerical rank rejects a tangent frame at machine roundoff without
         # imposing an application-level minimum angle of incidence.
         # Bad entries already have a finite-value error above; keep the rank
         # solver well-defined until the combined error flags reach the host.
         frame = torch.where(torch.isfinite(frame), frame, 0.)
-        checks.append((torch.any(torch.linalg.matrix_rank(frame) < 3),
-                       "parallel3d rays must not be tangent to the detector plane"))
+        if frame.dtype == torch.float64:
+            checks.append((torch.any(torch.linalg.matrix_rank(frame) < 3),
+                           "parallel3d rays must not be tangent to the detector plane"))
+        else:
+            # |det| via the triple product: matrix_rank has no MPS kernel. Its
+            # float32 threshold sits far above the FP64 rank tolerance, so it
+            # rejects a wider band of near-tangent frames than the FP64 check.
+            determinant = torch.sum(frame[..., 0] * torch.linalg.cross(frame[..., 1], frame[..., 2], dim=-1), dim=-1)
+            scale = torch.linalg.vector_norm(frame, dim=-2).prod(dim=-1)
+            checks.append((torch.any(determinant.abs() <= torch.finfo(determinant.dtype).eps * scale),
+                           "parallel3d rays must not be tangent to the detector plane"))
     if beam in ("fan", "cone") and not detector_surface:
-        # float64 avoids overflow of half-precision input and dtype mismatches in cross().
-        source = geometry[0].to(device=device, dtype=torch.float64)
-        principal = geometry[1].to(source.device, torch.float64) - source
+        source = _cast_work(geometry[0], device)
+        principal = _cast_work(geometry[1], source.device) - source
         lengths = torch.linalg.vector_norm(principal, dim=1)
         checks.append((torch.any(lengths == 0), "source and detector center must differ in every view"))
         principal = principal / lengths[:, None]
-        det_u = geometry[2].to(source.device, torch.float64)
+        det_u = _cast_work(geometry[2], source.device)
         if beam == "fan":
             facing = principal[:, 0] * det_u[:, 1] - principal[:, 1] * det_u[:, 0]
         else:
-            normal = torch.linalg.cross(det_u, geometry[3].to(source.device, torch.float64), dim=1)
+            normal = torch.linalg.cross(det_u, _cast_work(geometry[3], source.device), dim=1)
             facing = torch.sum(principal * normal, dim=1)
         checks.append((torch.any(torch.abs(facing) < 1e-4),
                        "the detector must not be edge-on to the source in any view"))
@@ -442,6 +450,7 @@ class Projector:
         self.native_volume_layout = bool(native_volume_layout)
         self.precision = precision
         self.dtype = torch.float64 if precision == "float64" else torch.float32
+        self._validate_precision(trajectory)
         self._reference = None
         self.partition = partition
         if not isinstance(partition, str) or partition not in ("views", "space", "hybrid"):
@@ -571,6 +580,12 @@ class Projector:
                     UserWarning, stacklevel=2,
                 )
 
+    def _validate_precision(self, values):
+        if self.dtype == torch.float64 and any(
+                torch.device(getattr(value, "device", "cpu")).type == "mps"
+                for value in values):
+            raise ValueError("precision must be 'fast' or 'float32' on MPS; MPS provides no float64")
+
     def _validate_geometry_range(self):
         if self.dtype == torch.float32 and self.beam in ("fan", "cone") and self.detector_surface is None:
             # Both FP32 backends set up each segment from the endpoint nearer
@@ -578,10 +593,11 @@ class Projector:
             spacing = voxel_spacing_tuple(self.voxel_spacing, len(self.global_volume_shape))
             device = next((value.device for value in self._trajectory[:2] if value.is_cuda),
                           self._trajectory[0].device)
-            pitches = torch.tensor(spacing[::-1], device=device, dtype=torch.float64)
+            work = _work_dtype(device)
+            pitches = torch.tensor(spacing[::-1], device=device, dtype=work)
             distances = torch.stack([
                 torch.linalg.vector_norm(
-                    component.to(device=device, dtype=torch.float64) / pitches,
+                    component.to(device=device, dtype=work) / pitches,
                     dim=1,
                 )
                 for component in self._trajectory[:2]
@@ -2602,8 +2618,10 @@ class Projector:
         override = trajectory is not None
         if trajectory is None and self._learnable:
             trajectory = self._learnable
-        if trajectory is not None:
-            try:
+        try:
+            self._validate_precision(args)
+            if trajectory is not None:
+                self._validate_precision(trajectory)
                 _reject_trajectory_tangents(trajectory)
                 frames, views = _validate_trajectory(
                     trajectory, self.beam, detector_surface=self.detector_surface is not None,
@@ -2686,8 +2704,8 @@ class Projector:
                                 (coincident, "source and detector surface pixels must differ"),
                             ])
                             del staged
-            except (TypeError, ValueError, RuntimeError) as failure:
-                error = failure
+        except (TypeError, ValueError, RuntimeError) as failure:
+            error = failure
         # None must participate too: mixed override/default calls or gradient
         # layouts cannot skip a collective on only one rank.
         # Metadata transport needs ordinary storage-backed tensors, even when

diffct/reconstruction.py

@@ -350,8 +350,20 @@ def helical_view_weights(angles, source_z, slice_z):
     if not torch.is_floating_point(angles) or angles.ndim != 1 or len(angles) < 2:
         raise ValueError("angles must be a real vector with at least two views")
     output_dtype = angles.dtype
-    source_z = torch.as_tensor(source_z, device=angles.device, dtype=torch.float64)
-    slice_z = torch.as_tensor(slice_z, device=angles.device, dtype=torch.float64)
+    output_device = angles.device
+    if output_device.type == "mps":
+        # MPS has no float64: the scheduling math runs on the host. List
+        # inputs are created directly as float64 — converting through the
+        # default dtype loses precision for large coordinates, and a fused
+        # MPS->CPU float64 cast returns zeros on torch 2.14.
+        angles = angles.to(device="cpu").to(dtype=torch.float64)
+        source_z = (source_z.to(device="cpu") if isinstance(source_z, torch.Tensor)
+                    else torch.as_tensor(source_z, dtype=torch.float64, device="cpu")).to(dtype=torch.float64)
+        slice_z = (slice_z.to(device="cpu") if isinstance(slice_z, torch.Tensor)
+                   else torch.as_tensor(slice_z, dtype=torch.float64, device="cpu")).to(dtype=torch.float64)
+    else:
+        source_z = torch.as_tensor(source_z, device=output_device, dtype=torch.float64)
+        slice_z = torch.as_tensor(slice_z, device=output_device, dtype=torch.float64)
     if (source_z.shape != angles.shape or slice_z.ndim != 1 or len(slice_z) == 0
             or not all(torch.isfinite(t).all() for t in (angles, source_z, slice_z))):
         raise ValueError("angles/source_z must match and slice_z must be a finite vector")
@@ -385,7 +397,7 @@ def helical_view_weights(angles, source_z, slice_z):
                - torch.maximum(edges[:-1][None, :], lower[:, None])).clamp_min(0) * 0.5
     weights = torch.empty_like(overlap)
     weights[:, order] = overlap
-    return weights.to(output_dtype)
+    return weights.to(device=output_device, dtype=output_dtype)
 
 
 def _axial_helix(geometry, isocenter):

diffct/streaming.py

@@ -96,6 +96,13 @@ def _tensor_internal_overlap(tensor):
 def _tensor_overlap(first, second):
     if first.device != second.device:
         return False
+    if first.device.type == "mps":
+        # MPS tensors are allocator-backed and cannot wrap external memory,
+        # so distinct storages are distinct allocations. Cross-allocation
+        # data_ptr differences are not meaningful on MPS, so only views of
+        # one storage reach the stride solver, where offsets stay valid.
+        if first.untyped_storage() is not second.untyped_storage():
+            return False
     a, b = _tensor_span(first), _tensor_span(second)
     if a is None or b is None or max(a[0], b[0]) >= min(a[1], b[1]):
         return False

diffct/surfaces.py

@@ -4,7 +4,7 @@ from dataclasses import dataclass
 
 import torch
 
-from .utils import _raise_geometry_errors
+from .utils import _raise_geometry_errors, _work_dtype
 
 @dataclass(frozen=True)
 class _GeometrySpec:
@@ -96,8 +96,11 @@ def _sample_geometry(beam, detector_shape, spacing, geometry, spec, views, pixel
         dtype = torch.float32
         for component in (*frames, offsets):
             dtype = torch.promote_types(dtype, component.dtype)
-    elif dtype not in (torch.float32, torch.float64):
+    if dtype not in (torch.float32, torch.float64):
         raise ValueError("surface geometry dtype must be float32 or float64")
+    if dtype.itemsize > _work_dtype(device).itemsize:
+        # CPU-side FP64 geometry must not follow FP64-less compute devices.
+        dtype = _work_dtype(device)
     offset_dtype = offsets.dtype if native_legacy else dtype
     if beam in ("cone", "parallel3d"):
         offsets = offsets.to(device=device, dtype=offset_dtype)
@@ -139,13 +142,14 @@ def _sample_geometry(beam, detector_shape, spacing, geometry, spec, views, pixel
             endpoints = points.reshape(view_count, -1, source.shape[-1])
             checks.append((torch.any(torch.all(sources.float() == endpoints.float(), dim=-1)),
                            "source and detector surface pixels must differ"))
+            work = _work_dtype(sources.device)
             if isinstance(voxel_spacing, (tuple, list)):
-                pitches = sources.new_tensor(voxel_spacing[::-1], dtype=torch.float64)
-                source_distance = torch.linalg.vector_norm(sources.double() / pitches, dim=-1)
-                pixel_distance = torch.linalg.vector_norm(endpoints.double() / pitches, dim=-1)
+                pitches = sources.new_tensor(voxel_spacing[::-1], dtype=work)
+                source_distance = torch.linalg.vector_norm(sources.to(dtype=work) / pitches, dim=-1)
+                pixel_distance = torch.linalg.vector_norm(endpoints.to(dtype=work) / pitches, dim=-1)
             else:
-                source_distance = torch.linalg.vector_norm(sources.double(), dim=-1) / voxel_spacing
-                pixel_distance = torch.linalg.vector_norm(endpoints.double(), dim=-1) / voxel_spacing
+                source_distance = torch.linalg.vector_norm(sources.to(dtype=work), dim=-1) / voxel_spacing
+                pixel_distance = torch.linalg.vector_norm(endpoints.to(dtype=work), dim=-1) / voxel_spacing
             checks.append((torch.any(torch.minimum(source_distance, pixel_distance) > 1e6),
                            "surface source or pixel must lie within 1e6 voxels"))
             checks.append((torch.any(source_distance > 1e15) | torch.any(pixel_distance > 1e15),

diffct/utils.py

@@ -18,6 +18,26 @@ import torch
 from .constants import _DTYPE, _TPB_2D, _TPB_3D, _INF, _EPSILON
 
 
+def _work_dtype(device):
+    """Interval, validation and composition math runs in float64 where the device supports it.
+
+    The MPS backend has no float64, so those paths fall back to float32 there.
+    """
+    return torch.float32 if torch.device(device).type == "mps" else torch.float64
+
+
+def _cast_work(value, device):
+    """Cast value to the working dtype on device.
+
+    torch 2.14 MPS returns zeros for a fused MPS->CPU float64 cast, so those
+    conversions go through the source device first. Autograd stays connected.
+    """
+    dtype = _work_dtype(device)
+    if value.device.type == "mps" and dtype == torch.float64 and value.device != torch.device(device):
+        return value.cpu().to(dtype=dtype).to(device=device)
+    return value.to(device=device, dtype=dtype)
+
+
 def _raise_geometry_errors(checks):
     """Read one validity scalar; fetch individual error flags only on failure."""
     device = next((flag.device for flag, _ in checks if flag.is_cuda), checks[0][0].device)

{diffct-2.4.4.dist-info → diffct-2.5.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: diffct
-Version: 2.4.4
+Version: 2.5.0
 Summary: Differentiable CT projection and reconstruction: precision modes, streaming, distributed execution and geometry gradients
 Project-URL: Homepage, https://github.com/sypsyp97/diffct
 Project-URL: Documentation, https://sypsyp97.github.io/diffct/
@@ -10,7 +10,7 @@ Author-email: Yipeng Sun <yipeng.sun@fau.de>
 License-Expression: Apache-2.0
 License-File: LICENSE
 License-File: NOTICE
-Keywords: CT reconstruction,CUDA,PyTorch,ROCm,computed tomography,differentiable programming
+Keywords: Apple Silicon,CT reconstruction,CUDA,MPS,PyTorch,ROCm,computed tomography,differentiable programming
 Classifier: Development Status :: 5 - Production/Stable
 Classifier: Environment :: GPU :: NVIDIA CUDA
 Classifier: Intended Audience :: Science/Research
@@ -34,7 +34,8 @@ Description-Content-Type: text/markdown
 <p align="center">English · <a href="https://github.com/sypsyp97/diffct/blob/main/README.zh.md">简体中文</a></p>
 
 <p align="center">
-  Differentiable CUDA projectors for CT: arbitrary trajectories, multi-GPU, multi-node, geometry gradients.
+  Differentiable CT projectors: arbitrary trajectories, multi-GPU, multi-node, geometry gradients.<br>
+  Native NVIDIA CUDA kernels, plus a portable PyTorch backend for CPU, ROCm and Apple Silicon (MPS).
 </p>
 
 <p align="center">
@@ -66,7 +67,7 @@ Description-Content-Type: text/markdown
 </p>
 
 > **Version 2.0** replaces the circular-orbit API of 1.x with per-view trajectories. Code written for 1.x needs changes; see the [migration notes](https://github.com/sypsyp97/diffct/blob/main/docs/MIGRATION.md).
-> For Apple Silicon, use the [MLX port by Linda-Sophie Schneider](https://github.com/Linda-SophieSchneider/DiffCT-MLX).
+> Apple Silicon runs through the built-in MPS support (`backend="torch"`); a community [MLX port by Linda-Sophie Schneider](https://github.com/Linda-SophieSchneider/DiffCT-MLX) also exists.
 
 This README describes the source checkout. Features listed under [Unreleased](CHANGELOG.md#unreleased) require installation from source.
 
@@ -79,7 +80,7 @@ To our knowledge, diffct is the only open-source GPU CT library that combines ar
 - **Matched operators.** `project()` and `backproject()` form an exact adjoint pair for the cell-constant Siddon model. Both support PyTorch autograd, with volume/sinogram gradients and Hessian-vector products.
 - **Large volumes.** Automatic spatial, view and pixel blocks with bounded transfer buffers control working memory. `project_into`/`backproject_into` accept caller-owned tensor or disk stores; tensor autograd keeps complete data on its input device. See [chunked execution](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#chunked-execution).
 - **Multi-GPU and multi-node.** Use `devices=[0, 1, 2, 3]` in one process, or torchrun and NCCL across ranks. Views, space and hybrid partitions control replication of volumes and projections. Speedup depends on workload and communication costs.
-- **Precision and training.** Explicit native/PyTorch backends, FP32/FP64 modes, per-axis voxel spacing, leading batch/channel dimensions and data `torch.func` transforms. The portable path supports CPU and PyTorch CUDA/ROCm devices; see [execution contracts](docs/EXECUTION.md).
+- **Precision and training.** Explicit native/PyTorch backends, FP32/FP64 modes, per-axis voxel spacing, leading batch/channel dimensions and data `torch.func` transforms. The portable path supports CPU, PyTorch CUDA/ROCm devices and Apple Silicon GPUs (MPS, `backend="torch"`); see [execution contracts](docs/EXECUTION.md).
 - **Solver API.** `cgls`, `sirt` and `tv_reconstruction` report residual history, tolerances and callbacks, with scan/data fingerprints and complete checkpoint state. `cgls_into` keeps state in bounded mapped blocks.
 - **Differentiable FBP and FDK.** `fbp`, `fdk` and approximate `helical_fdk` compose the public filters, weights and analytical backprojectors. Their backward kernels transpose the actual interpolation stencil; custom filters and weights can be learned. See the [analytical guide](docs/source/analytical_learning.rst).
 
@@ -95,7 +96,7 @@ pip install "diffct[cu12]"    # PyTorch built for CUDA 12
 pip install "diffct[cu13]"    # PyTorch built for CUDA 13
 ```
 
-The extra installs the CUDA compiler libraries Numba needs. For CPU or PyTorch ROCm execution, install the appropriate PyTorch build and plain `pip install diffct`, then select `backend="torch"`. Importing diffct and using this backend require no CUDA runtime; Numba CUDA initializes when native execution is used. Plain installation also works for native execution when a compatible CUDA Toolkit is already present.
+The extra installs the CUDA compiler libraries Numba needs. For CPU, PyTorch ROCm or Apple Silicon (MPS) execution, install the appropriate PyTorch build and plain `pip install diffct`, then select `backend="torch"`. Importing diffct and using this backend require no CUDA runtime; Numba CUDA initializes when native execution is used. Plain installation also works for native execution when a compatible CUDA Toolkit is already present.
 
 From source, with the examples:
 

{diffct-2.4.4.dist-info → diffct-2.5.0.dist-info}/RECORD

@@ -1,24 +1,24 @@
-diffct/__init__.py,sha256=Euic5adzAfIdFslOxF_81MyphkyZK2sPT0ETeYofrRY,2965
-diffct/_solver_state.py,sha256=ExWTTGftGOsPo9MQwxv-ylmYvhiPYXUo3Tmhi6VEvg0,16606
+diffct/__init__.py,sha256=jZeEhRW9GpQpRVrxV2LbuyBs33b1RzoIT7cqJE2pQVA,2965
+diffct/_solver_state.py,sha256=qzyXgJCx48gWXPcJDeANXuAhAZopDzvK7ByckABVN9I,18238
 diffct/_solver_tv.py,sha256=2x16Bq5Bx9WCBszgPy3K-wicsH9AeuzzJS82uyiHBkE,5847
 diffct/analytical.py,sha256=LGYA5NEXVJ0qKwOWPXIZMXY2EOKvLto_w7bBcep5N2I,33398
-diffct/analytical_execution.py,sha256=sS3WFUyaaBJnr8d3XiUJoWTxQujoiB3-oPbXPvZcL9c,27069
-diffct/backends.py,sha256=pV2a3ODyhZWEV0nPAI82yPqNDlT85btKnTbCag4pCtQ,70473
+diffct/analytical_execution.py,sha256=Tc4x0IglgUdQ7IGoZu2Jt_5onCGibX8XwWx7lzhXSLc,27248
+diffct/backends.py,sha256=oQbNPJY-zyUBzG9uw384duF66Oo6UuiVfi3wBqFISd0,71906
 diffct/chunking.py,sha256=sdZHajLpHEh1mq4OHCAQZzU72QGQrMRN5Ft0bxE3nTM,11480
 diffct/constants.py,sha256=wgOfwUZW8CZsj0cj0f6gBLCxbq2qa4xY47QiPEUvWrc,3747
 diffct/differentiable.py,sha256=cFirhqXRpUemL7y7Sl_0uqZKGJNj__ptQEcTwYNKzys,1879
 diffct/distribution.py,sha256=iLC3BgAQ9lOkndbEOvNB_j_UaA2Ug9w1KiTUHkOzG8o,9982
-diffct/geometry.py,sha256=iwWXQksIJGOPSuNWJ8tImWj6yPZCC71G3RBh2VDVylY,42581
-diffct/operators.py,sha256=aSHd1sIJ1_fc4epMFS8jW183r9dI206cvbB8YsAtNjs,170540
+diffct/geometry.py,sha256=1oaXkwm3MoIgCcBP1W38MhVikhz-bH9IeWRyIitdI0s,44699
+diffct/operators.py,sha256=yWvITQcAKtGaUlY94xSA65ItGdynmv-jj2PmgdQExBg,171550
 diffct/pipeline.py,sha256=jp7AvDeQQOME3OPkO9XTcjWFLvOLrv8wczhAf1jcklM,10891
 diffct/projectors.py,sha256=fYfjw-BU1odP0htmlgVWdJk5pL0zP6MfgArOzdGCyq0,45900
-diffct/reconstruction.py,sha256=z_5dPffzs3PEaLl8rdV_l2FpilSNkUqheWgEEfywJM8,29564
+diffct/reconstruction.py,sha256=8RTbtouesjFW4CkcuzZ-wDinfeupQZnEuiqveMR0yF8,30415
 diffct/solvers.py,sha256=B7M4akFUTuwb_MBpQbpZLdlMv5Eyxy3WczExDAs613s,18000
 diffct/storage.py,sha256=EdcFTzRtEAPRvAZhRPrSf8gRNF5o_JVbPLE8eKPhNdY,4153
-diffct/streaming.py,sha256=2LB5xpl6vIJE1kXMBtQbG0OVxlutEEqXLBrcvlzBHl8,6086
+diffct/streaming.py,sha256=5Iu4zY9bhP5ALw3mqvhHe-VbM1eesyAGzYRnZWXRp8w,6512
 diffct/surface_gradients.py,sha256=LeM_eDY1Alv604SdiPUx2CWFzr-J9NqCEcJ_ovVFGa8,6290
-diffct/surfaces.py,sha256=-tp7tJi2GVy1lE4-Yi46hVQ6BMMF08Pfck3rkibgF-Y,8370
-diffct/utils.py,sha256=ZOZpR_qpVSxv_uSQJ6--x-RXAqUrZJjWXYJoPKhu5nY,15240
+diffct/surfaces.py,sha256=p7vk-W0KR7pbJcb_zJqQxNG8-iocmBkCe9WAq6H9ac8,8609
+diffct/utils.py,sha256=1Bn37Ay9RsqgofIYNYZilpjT-jkJWv_GBF2Q-vIXqz8,16023
 diffct/kernels/__init__.py,sha256=YNdsqAkexN25pI7QaHy4phAnQNICQ7tO8K8S1zt7TnI,1899
 diffct/kernels/analytical_transpose.py,sha256=3aiYjMm4IVJUZ9DJaHZLgA_cPgoTc8pvoZUJ6KZCYtI,10006
 diffct/kernels/cone_beam.py,sha256=EQRmd5Q5MGFrtk2iddIgxV4xlyxhjQh-_3u3ZTa1Z0Q,30670
@@ -27,8 +27,8 @@ diffct/kernels/geometry_vjp.py,sha256=JZ5At_0EtsCBacxt38hrxGG43USEA_4eiqz0PHF1zM
 diffct/kernels/parallel_beam.py,sha256=M6Hjf8B66MNpb6ze5ScYEZNQ151tPd0TODrMFs8YMV8,18764
 diffct/kernels/spacing.py,sha256=iRVlHZpLi15dS1sjzNJMohfxooKjschF6Y1wTDe7SZA,2031
 diffct/kernels/surface_geometry.py,sha256=tTJXykpK7lOfdPr3Oz1bCCFm7pRnbttGJHTkHuAxRIA,6322
-diffct-2.4.4.dist-info/METADATA,sha256=xmnyKp7ix75ZC1SZkP0xC8g6_0Ztf_wab50ifFQOgHU,17865
-diffct-2.4.4.dist-info/WHEEL,sha256=W3fkpkm7-wf9vBI5Z-7s0eWkeM-spu78I8Neb98DeEg,87
-diffct-2.4.4.dist-info/licenses/LICENSE,sha256=a1fjzGs44LpDnV_LmA2Tzp8z-eJKoNVVdLOnh-F6J6E,11344
-diffct-2.4.4.dist-info/licenses/NOTICE,sha256=Qb3qWsJ3TMqipmMm2_jS8Ryrq9XyHtdjc3qH1q8ELi4,798
-diffct-2.4.4.dist-info/RECORD,,
+diffct-2.5.0.dist-info/METADATA,sha256=1a01Ijl6XVM-8m0uUwySgpJKLNqIZGdj-BBhXnvvrK4,18115
+diffct-2.5.0.dist-info/WHEEL,sha256=W3fkpkm7-wf9vBI5Z-7s0eWkeM-spu78I8Neb98DeEg,87
+diffct-2.5.0.dist-info/licenses/LICENSE,sha256=a1fjzGs44LpDnV_LmA2Tzp8z-eJKoNVVdLOnh-F6J6E,11344
+diffct-2.5.0.dist-info/licenses/NOTICE,sha256=Qb3qWsJ3TMqipmMm2_jS8Ryrq9XyHtdjc3qH1q8ELi4,798
+diffct-2.5.0.dist-info/RECORD,,