segresnet-mlx 0.2.1__tar.gz → 0.2.3__tar.gz

{segresnet_mlx-0.2.1 → segresnet_mlx-0.2.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: segresnet-mlx
-Version: 0.2.1
+Version: 0.2.3
 Summary: Native MLX SegResNet with NumPy sliding-window logit inference
 Project-URL: Homepage, https://github.com/MedMLX/segresnet-mlx
 Project-URL: Documentation, https://github.com/MedMLX/segresnet-mlx/blob/main/docs/technical.md
@@ -12,7 +12,7 @@ License-File: LICENSE
 License-File: THIRD_PARTY_NOTICES.md
 Classifier: Operating System :: MacOS :: MacOS X
 Requires-Python: <3.15,>=3.13
-Requires-Dist: medmlx-core==0.2.3
+Requires-Dist: medmlx-core==0.2.4
 Requires-Dist: nibabel<6,>=5.3
 Requires-Dist: numpy<3,>=2.0
 Provides-Extra: conversion
@@ -47,26 +47,22 @@ Segment brain tumors, renal structures and whole-body CT anatomy with MONAI SegR
 ## Quick start
 
 ```bash
-uv venv --python 3.13 --seed
+uv venv --python 3.13
 source .venv/bin/activate
-python -m pip install "segresnet-mlx==0.2.1" huggingface-hub
+uv pip install segresnet-mlx huggingface-hub
 ```
 
-The tagged GitHub release is an alternative installation route:
-
-```bash
-repo="https://github.com/MedMLX/segresnet-mlx.git@v0.2.1"
-python -m pip install "segresnet-mlx @ git+$repo" huggingface-hub
-```
-
-Download the public converted BraTS bundle separately, then run:
+Download the converted [BraTS](https://huggingface.co/MedMLX/segresnet-brats-mlx) weights, then run:
 
 ```python
 from huggingface_hub import hf_hub_download
 from segresnet_mlx.medmlx import run
 
 bundle = "brats_mri_segmentation"
-weights = hf_hub_download("MedMLX/segresnet-brats-mlx", "segresnet-brats.safetensors")
+weights = hf_hub_download(
+    "MedMLX/segresnet-brats-mlx", "segresnet-brats.safetensors",
+    revision="14f9cda3649965729663c7365c49fd839c81a01b",
+)
 
 result = run(
     bundle=bundle,
@@ -80,12 +76,9 @@ result = run(
 print(result["outputs"]["mask"])
 ```
 
-Keep input volumes on the same voxel grid. The package contains code only.
-Public converted bundles are available separately for
-[BraTS](https://huggingface.co/MedMLX/segresnet-brats-mlx) and
-[renal CECT](https://huggingface.co/MedMLX/segresnet-renal-cect-mlx); follow each
-weight card and its terms. Docs cover local upstream checkpoint conversion and
-whole-body CT bundles. Qualification differs by bundle; see the technical notes.
+Keep input volumes on the same voxel grid. Converted weights are also available for
+[renal CECT](https://huggingface.co/MedMLX/segresnet-renal-cect-mlx). Docs cover
+checkpoint conversion and whole-body CT.
 
 ## License
 

{segresnet_mlx-0.2.1 → segresnet_mlx-0.2.3}/README.md

@@ -18,26 +18,22 @@ Segment brain tumors, renal structures and whole-body CT anatomy with MONAI SegR
 ## Quick start
 
 ```bash
-uv venv --python 3.13 --seed
+uv venv --python 3.13
 source .venv/bin/activate
-python -m pip install "segresnet-mlx==0.2.1" huggingface-hub
+uv pip install segresnet-mlx huggingface-hub
 ```
 
-The tagged GitHub release is an alternative installation route:
-
-```bash
-repo="https://github.com/MedMLX/segresnet-mlx.git@v0.2.1"
-python -m pip install "segresnet-mlx @ git+$repo" huggingface-hub
-```
-
-Download the public converted BraTS bundle separately, then run:
+Download the converted [BraTS](https://huggingface.co/MedMLX/segresnet-brats-mlx) weights, then run:
 
 ```python
 from huggingface_hub import hf_hub_download
 from segresnet_mlx.medmlx import run
 
 bundle = "brats_mri_segmentation"
-weights = hf_hub_download("MedMLX/segresnet-brats-mlx", "segresnet-brats.safetensors")
+weights = hf_hub_download(
+    "MedMLX/segresnet-brats-mlx", "segresnet-brats.safetensors",
+    revision="14f9cda3649965729663c7365c49fd839c81a01b",
+)
 
 result = run(
     bundle=bundle,
@@ -51,12 +47,9 @@ result = run(
 print(result["outputs"]["mask"])
 ```
 
-Keep input volumes on the same voxel grid. The package contains code only.
-Public converted bundles are available separately for
-[BraTS](https://huggingface.co/MedMLX/segresnet-brats-mlx) and
-[renal CECT](https://huggingface.co/MedMLX/segresnet-renal-cect-mlx); follow each
-weight card and its terms. Docs cover local upstream checkpoint conversion and
-whole-body CT bundles. Qualification differs by bundle; see the technical notes.
+Keep input volumes on the same voxel grid. Converted weights are also available for
+[renal CECT](https://huggingface.co/MedMLX/segresnet-renal-cect-mlx). Docs cover
+checkpoint conversion and whole-body CT.
 
 ## License
 

{segresnet_mlx-0.2.1 → segresnet_mlx-0.2.3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "segresnet-mlx"
-version = "0.2.1"
+version = "0.2.3"
 description = "Native MLX SegResNet with NumPy sliding-window logit inference"
 readme = "README.md"
 requires-python = ">=3.13,<3.15"
@@ -13,7 +13,7 @@ license = "Apache-2.0"
 license-files = ["LICENSE", "THIRD_PARTY_NOTICES.md"]
 authors = [{ name = "Joseph Sandoval" }, { name = "Alfredo Sandoval Jr." }]
 dependencies = [
-  "medmlx-core==0.2.3",
+  "medmlx-core==0.2.4",
   "nibabel>=5.3,<6",
   "numpy>=2.0,<3",
 ]
@@ -78,7 +78,7 @@ ban-relative-imports = "all"
 "typing_extensions.Any" = {msg = "Use a precise type and parse untyped values at the boundary."}
 
 [tool.uv.sources]
-medmlx-core = { git = "https://github.com/MedMLX/medmlx.git", rev = "e765f0a94aa31e8dfcdc599e4a4206eaca84f8b0", subdirectory = "packages/core" }
+medmlx-core = { git = "https://github.com/MedMLX/medmlx.git", rev = "f98e0924f50228e2b9509474947b9672054b3217", subdirectory = "packages/core" }
 
 [tool.hatch.build.targets.sdist]
 only-include = [

{segresnet_mlx-0.2.1 → segresnet_mlx-0.2.3}/src/segresnet_mlx/__init__.py

@@ -21,7 +21,7 @@ if TYPE_CHECKING:
     from segresnet_mlx.weights import load_model as load_model
     from segresnet_mlx.weights import required_shapes as required_shapes
 
-__version__ = "0.2.1"
+__version__ = "0.2.3"
 
 _EXPORTS = {
     "BUNDLE_CONFIGS": "segresnet_mlx.predictor",

{segresnet_mlx-0.2.1 → segresnet_mlx-0.2.3}/src/segresnet_mlx/_types.py

@@ -47,8 +47,9 @@ if TYPE_CHECKING:
             overlap: float,
             mode: str,
             padding_mode: str,
-            with_coord: bool,
+            with_coord: bool = False,
             predictor_input: Literal["numpy"],
+            preserve_window_strides: bool = False,
         ) -> Float32Array: ...
 
     class CoreSlidingWindow(Protocol):

{segresnet_mlx-0.2.1 → segresnet_mlx-0.2.3}/src/segresnet_mlx/medmlx.py

@@ -292,16 +292,23 @@ def run(
     ``brats_mri_segmentation`` (``t1c_path``, ``t1_path``, ``t2_path``,
     ``flair_path``: co-registered MRI on one grid), or ``wholeBody_ct_segmentation``
     / ``wholeBody_ct_segmentation_lowres`` (``image_path``: one CT volume, segmented
-    at 1.5 mm or 3 mm respectively). ``weights_path`` is the bundle's
-    ``models/model.pt`` (or ``models/model_lowres.pt`` for the lowres whole-body
-    variant) converted with :func:`segresnet_mlx.convert_checkpoint`.
+    at 1.5 mm or 3 mm respectively). ``weights_path`` is a native ``.safetensors``
+    file: ``segresnet-brats.safetensors`` from ``MedMLX/segresnet-brats-mlx`` for
+    BraTS, or ``segresnet-renal-cect.safetensors`` from
+    ``MedMLX/segresnet-renal-cect-mlx`` for renal CECT. These published files need
+    no conversion. For other checkpoints, convert the original ``models/model.pt``
+    (``models/model_lowres.pt`` for lowres whole-body) with
+    :func:`segresnet_mlx.convert_checkpoint` and pass the resulting file.
 
     Inputs must be finite 3D NIfTI volumes; multi-input bundles need identical shape
     and affine (atol 1e-5). Writes ``<bundle>_mask.nii.gz`` in ``output_dir`` on the
     first input's shape and reader-selected affine (see ``Volume``): renal labels 0-5
     and whole-body labels 0-104 as float32, BraTS labels 0, 1, 2, 4 as uint8, as the
     bundles' ``SaveImaged`` wrote them. An existing mask is replaced only when
-    ``overwrite`` is true; output aliases of input images are always refused.
+    ``overwrite`` is true; output aliases of input images are refused by a
+    preflight check. Spatial affine values are assumed to be in millimeters; the
+    NIfTI spatial-unit field is not checked or converted. Paths are not locked
+    against concurrent replacement after the alias check.
 
     ``arithmetic_profile="compensated-fp32"`` uses Metal float pairs for
     convolution and GroupNorm accumulation, with float32 layer boundaries.
@@ -312,6 +319,27 @@ def run(
     requires transpose-convolution upsampling and compares with the unchanged
     original graph evaluated in float64 throughout its network and inferer.
 
+    Args:
+        bundle: ``brats_mri_segmentation``, ``renalStructures_CECT_segmentation``,
+            ``wholeBody_ct_segmentation`` or ``wholeBody_ct_segmentation_lowres``.
+        weights_path: Path to the matching native ``.safetensors`` file, such as
+            ``segresnet-brats.safetensors`` or ``segresnet-renal-cect.safetensors``.
+            Hugging Face cache file paths are accepted; directories and Torch ``.pt``
+            checkpoints are not.
+        output_dir: Directory for ``<bundle>_mask.nii.gz``; created if missing.
+        arterial_path: Renal arterial-phase 3D NIfTI; required for the renal bundle.
+        venous_path: Renal venous-phase 3D NIfTI on the arterial grid; required for renal.
+        excretory_path: Renal excretory-phase 3D NIfTI on the arterial grid; required
+            for renal.
+        t1c_path: BraTS contrast-enhanced T1 3D NIfTI; required for BraTS.
+        t1_path: BraTS T1 3D NIfTI on the T1c grid; required for BraTS.
+        t2_path: BraTS T2 3D NIfTI on the T1c grid; required for BraTS.
+        flair_path: BraTS FLAIR 3D NIfTI on the T1c grid; required for BraTS.
+        image_path: Single 3D CT NIfTI; required for either whole-body bundle.
+        overwrite: Replace an existing output mask when true; defaults to false.
+        arithmetic_profile: ``torch-order-fp32`` (default), ``compensated-fp32``,
+            or ``compensated-fp48``; the last supports only the renal bundle.
+
     Returns:
         ``{"outputs": {"mask": path}, "details": {...}}``; ``details`` is described
         by ``segresnet_mlx._records.Details`` and always has
@@ -534,7 +562,9 @@ def postprocess_wholebody(logits: Float32Array, inverse: SpatialInverse) -> Floa
 
 
 def _weights_file(weights_path: str) -> Path:
-    path = Path(weights_path).expanduser().resolve()
+    # MLX identifies the file format from the supplied name. Hub cache paths are
+    # named .safetensors symlinks whose underlying content-addressed blobs have no suffix.
+    path = Path(weights_path).expanduser().absolute()
     hint = (
         "Convert the bundle's models/model.pt once with "
         "segresnet_mlx.convert_checkpoint(model_pt, destination, spec=...) "
@@ -556,8 +586,9 @@ def _weights_file(weights_path: str) -> Path:
 def load_inputs(roles: list[tuple[str, str | None]]) -> list[Volume]:
     """Load each ``(role, path)`` and require one shared voxel grid.
 
-    Every volume must match the first one's shape and affine (``atol=1e-5``); the
-    runner does not resample, so a mismatch raises ``InvalidInputError``.
+    Every volume must match the first one's shape and affine (``atol=1e-5``).
+    This check does not align mismatched inputs; bundle preprocessing may still
+    resample the admitted grid afterward.
     """
     volumes = [load_volume(Path(str(path)).expanduser().resolve(), role) for role, path in roles]
     reference = volumes[0]
@@ -573,7 +604,12 @@ def load_inputs(roles: list[tuple[str, str | None]]) -> list[Volume]:
 
 
 def load_volume(path: Path, role: str) -> Volume:
-    """Read finite 3D voxels using the affine selected by the upstream nibabel reader."""
+    """Read 3D voxels using the affine selected by the upstream nibabel reader.
+
+    Source voxels are checked for finiteness before conversion to float32; values
+    outside the float32 range can overflow during that conversion. Spatial units
+    and affine invertibility are not checked here.
+    """
     from segresnet_mlx._nifti import ImageFileError, nib
 
     if not path.is_file():

{segresnet_mlx-0.2.1 → segresnet_mlx-0.2.3}/src/segresnet_mlx/network.py

@@ -21,7 +21,8 @@ def group_mean(values: Array, mx: MLX) -> Array:
     """Mean over the last axis of ``(N, G, S)`` values, summed in 256-element tiles.
 
     Tiling bounds the float32 accumulation depth of a group that spans a whole
-    volume. Returns ``(N, G, 1)``. ``SegResNet`` does not call this (its GroupNorm
+    volume. The last dimension must be nonempty. Returns ``(N, G, 1)``. ``SegResNet``
+    does not call this (its GroupNorm
     comes from the ``medmlx_core`` operators); it is exercised by
     ``tests/test_equivalence.py``.
     """
@@ -34,16 +35,20 @@ def group_mean(values: Array, mx: MLX) -> Array:
 class SegResNet:
     """Raw-logit forward pass for MONAI's 3-D, group-normalized SegResNet.
 
-    Tensors are channels-first ``(N, C, D, H, W)`` float32 at every layer boundary.
+    Tensors are channels-first ``(N, C, D, H, W)``; the supported inference
+    contract uses float32 inputs and returned logits. This class checks rank,
+    channel count and spatial divisibility but does not itself check input dtype
+    or finiteness; the public patch and sliding-window helpers do.
     ``weights`` uses MONAI's state-dict keys and Torch parameter layouts
     (convolution weights are ``(out, in, kD, kH, kW)``); ``required_shapes`` in
     ``segresnet_mlx.weights`` lists them. Inference only: dropout is omitted and
     no activation is applied to the output.
 
-    ``arithmetic_profile`` selects the accumulation scheme (see the README):
+    ``arithmetic_profile`` selects the accumulation scheme (see ``docs/technical.md``):
     ``"torch-order-fp32"`` (default), ``"compensated-fp32"``, or
     ``"compensated-fp48"``, which needs ``upsample_mode="deconv"`` and raises
-    ``ValueError`` otherwise.
+    ``ValueError`` otherwise. The first two profiles use float32 layer boundaries;
+    ``compensated-fp48`` carries paired float32 values through the hidden layers.
     """
 
     def __init__(

{segresnet_mlx-0.2.1 → segresnet_mlx-0.2.3}/src/segresnet_mlx/predictor.py

@@ -6,7 +6,7 @@ import time
 from collections.abc import Callable
 from contextlib import nullcontext
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Literal, cast
+from typing import TYPE_CHECKING, cast
 
 from segresnet_mlx._records import BundleConfig, InfererConfig
 from segresnet_mlx._runtime import import_mlx
@@ -50,8 +50,9 @@ class SlidingWindowConfig:
     def from_bundle_config(cls, config: InfererConfig | BundleConfig) -> SlidingWindowConfig:
         """Read the staged JSON's inferer section without evaluating bundle expressions.
 
-        Missing keys fall back to MONAI's defaults. Values are taken literally, so a
-        reference such as ``@displayable_configs#sw_overlap`` is not resolved.
+        ``roi_size`` is required; the other listed keys default to MONAI's
+        values. Values are taken literally, so a reference such as
+        ``@displayable_configs#sw_overlap`` is not resolved.
         """
         inferer = cast(InfererConfig, config.get("inferer", config))
         return cls(
@@ -71,13 +72,6 @@ BUNDLE_CONFIGS = {
     for name, bundle in BUNDLES.items()
 }
 
-_PAD_MODES: dict[str, Literal["constant", "reflect", "edge", "wrap"]] = {
-    "constant": "constant",
-    "reflect": "reflect",
-    "replicate": "edge",
-    "circular": "wrap",
-}
-
 
 class BundlePredictor:
     """Synchronized MLX window callback with an injected network.
@@ -195,48 +189,18 @@ class BundlePredictor:
             require_ncdhw(host, name="SegResNet input")
             if self.telemetry is not None:
                 self.telemetry.set_metadata({"sw_batch_size": configuration.sw_batch_size})
-            callback: (
-                Callable[[HostArray | Array], Float32Array]
-                | Callable[[HostArray | Array, list[list[slice]]], Float32Array]
-            ) = self._forward_window
-            use_coordinates = configuration.sw_batch_size == 1
-            if use_coordinates:
-                # MONAI preserves the sliced input's strides for batch size 1;
-                # the shared core concatenates even a single patch. MLX selects
-                # different reduction kernels for strided/contiguous inputs.
-                # Open gap: a core change that skips the single-patch
-                # concatenation would make this workaround unnecessary.
-                padding = [(0, 0), (0, 0)]
-                for size, roi in zip(host.shape[2:], configuration.roi_size, strict=True):
-                    extra = max(roi - size, 0)
-                    padding.append((extra // 2, extra - extra // 2))
-                # Torch padding materializes a contiguous array; unpadded data
-                # retains its original strides. Match the bundle's padding_mode
-                # so this fast path agrees with the core's own padding.
-                pad_kind = _PAD_MODES[configuration.padding_mode]
-                padded = (
-                    np.ascontiguousarray(np.pad(host, padding, mode=pad_kind))
-                    if any(before or after for before, after in padding)
-                    else host
-                )
-
-                def coordinate_callback(
-                    values: HostArray | Array, coordinates: list[list[slice]]
-                ) -> Float32Array:
-                    return self._forward_window(padded[tuple(coordinates[0])])
-
-                callback = coordinate_callback
-
             return sliding_window_inference(
                 host,
                 configuration.roi_size,
                 configuration.sw_batch_size,
-                callback,
+                self._forward_window,
                 overlap=configuration.overlap,
                 mode=configuration.mode,
                 padding_mode=configuration.padding_mode,
-                with_coord=use_coordinates,
                 predictor_input="numpy",
+                # MLX reduction kernels depend on layout. Match MONAI's sliced
+                # batch-one input without a second padded volume or patch copy.
+                preserve_window_strides=True,
             )
         except BaseException as error:
             if self.telemetry is not None:
@@ -287,17 +251,21 @@ def predict_wide_logits(
         if not np.isfinite(output).all():
             raise RuntimeError("SegResNet produced non-finite paired logits")
         del result
-        mx.clear_cache()
         return output
 
-    logits = sliding_window_inference(
-        values.astype(np.float64),
-        configuration.roi_size,
-        configuration.sw_batch_size,
-        window,
-        overlap=configuration.overlap,
-        mode=configuration.mode,
-        padding_mode=configuration.padding_mode,
-        predictor_input="numpy",
-    )
-    return logits.astype(np.float32)
+    try:
+        logits = sliding_window_inference(
+            values.astype(np.float64),
+            configuration.roi_size,
+            configuration.sw_batch_size,
+            window,
+            overlap=configuration.overlap,
+            mode=configuration.mode,
+            padding_mode=configuration.padding_mode,
+            predictor_input="numpy",
+        )
+        return logits.astype(np.float32)
+    finally:
+        # Reuse same-shaped Metal allocations across windows, including when
+        # paired arithmetic requires a wider host accumulator.
+        mx.clear_cache()

{segresnet_mlx-0.2.1 → segresnet_mlx-0.2.3}/src/segresnet_mlx/topology.py

@@ -6,7 +6,7 @@ import math
 from dataclasses import dataclass
 from typing import Literal
 
-# Accumulation scheme for convolution and GroupNorm; see the README for the three
+# Accumulation scheme for convolution and GroupNorm; see docs/technical.md for the three
 # profiles. ``compensated-fp48`` additionally requires transpose-convolution upsampling.
 type ArithmeticProfile = Literal["torch-order-fp32", "compensated-fp32", "compensated-fp48"]
 

{segresnet_mlx-0.2.1 → segresnet_mlx-0.2.3}/src/segresnet_mlx/transforms.py

@@ -10,7 +10,8 @@ bundle's ``MergeClassesd`` and the NIfTI writer. Spatial resampling follows MONA
 as ``Spacing`` does and returned as float32.
 
 Image arrays are channel-first float32 ``(C, *spatial)`` and affines are 4x4 float64
-voxel-to-world matrices as stored in the NIfTI header; ``pixdim`` values are in mm.
+voxel-to-world matrices as stored in the NIfTI header. Bundle ``pixdim`` targets
+are in mm; these functions do not inspect or convert the input's spatial units.
 """
 
 from __future__ import annotations
@@ -373,7 +374,11 @@ def brats_labels(binary: Float32Array) -> npt.NDArray[np.int64]:
 
 
 def write_nifti(path: Path, data: HostArray, affine: Float64Array, dtype: npt.DTypeLike) -> None:
-    """MONAI ``NibabelWriter``: cast, then set sform and qform (code 1) from the affine."""
+    """Cast and write a NIfTI with sform and qform code 1.
+
+    The sform stores the affine within NIfTI precision. The qform cannot
+    represent shear, which NiBabel removes when encoding that form.
+    """
     array = np.asarray(data).astype(dtype, copy=False)
     matrix = to_affine_nd(3, affine)
     image = nib.Nifti1Image(array, matrix)
@@ -517,12 +522,15 @@ def _grid_sample(
     return result
 
 
-def _fma_array(left: Float64Array, right: float, addend: Float64Array) -> Float64Array:
+def _fma_array(
+    left: Float64Array, right: float | Float64Array, addend: Float64Array
+) -> Float64Array:
     """Elementwise ``left * right + addend`` with one rounding (error-free transforms)."""
-    left, addend = np.broadcast_arrays(np.asarray(left, np.float64), np.asarray(addend, np.float64))
-    right_value = np.float64(right)
+    left, right_value, addend = np.broadcast_arrays(
+        np.asarray(left, np.float64), np.asarray(right, np.float64), np.asarray(addend, np.float64)
+    )
     product = left * right_value
-    if right_value == 0 or not np.any(left):
+    if not np.any(right_value) or not np.any(left):
         return product + addend
     error = _product_error(left, right_value, product)
     total = product + addend
@@ -531,7 +539,9 @@ def _fma_array(left: Float64Array, right: float, addend: Float64Array) -> Float6
     return total + (sum_error + error)
 
 
-def _product_error(left: Float64Array, right: np.float64, product: Float64Array) -> Float64Array:
+def _product_error(
+    left: Float64Array, right: Float64Array, product: Float64Array
+) -> Float64Array:
     """Dekker's exact remainder ``left * right - product``."""
 
     def split[T: (Float64Array, np.float64)](value: T) -> tuple[T, T]:
@@ -580,7 +590,12 @@ def _sample_separable(
     gy: Float64Array,
     gz: Float64Array,
 ) -> Float64Array:
-    """Trilinear corners for an axis-aligned grid, in ATen's weight and sum order."""
+    """Trilinear corners in ATen's order, with its fused multiply-add accumulation.
+
+    Separate product/sum rounding can move a float64 interpolation across a
+    float32 midpoint. Preserve the source CPU kernel's fusion before the final
+    float32 cast; the network still receives ordinary float32 values.
+    """
     depth, height, width = data.shape[1:]
     x_index, x_weight = _corners(_source_index(gx, width))
     y_index, y_weight = _corners(_source_index(gy, height))
@@ -602,7 +617,7 @@ def _sample_separable(
             .astype(np.float64)
         )
         inside = ((zi < depth)[along_d] & (yi < height)[along_h]) & (xi < width)[along_w]
-        result += np.where(inside, values * weight, 0.0)
+        result = np.where(inside, _fma_array(values, weight, result), result)
     return result
 
 
@@ -612,7 +627,7 @@ def _sample(
     gy: Float64Array,
     gz: Float64Array,
 ) -> Float64Array:
-    """Trilinear corners for a general grid, in ATen's weight and sum order."""
+    """General-grid corners with the same fused accumulation as the separable path."""
     depth, height, width = data.shape[1:]
     x_index, x_weight = _corners(_source_index(gx, width))
     y_index, y_weight = _corners(_source_index(gy, height))
@@ -631,7 +646,7 @@ def _sample(
             np.minimum(yi, height - 1),
             np.minimum(xi, width - 1),
         ].astype(np.float64)
-        result += np.where(inside, values * weight, 0.0)
+        result = np.where(inside, _fma_array(values, weight, result), result)
     return result
 
 

{segresnet_mlx-0.2.1 → segresnet_mlx-0.2.3}/src/segresnet_mlx/weights.py

@@ -78,6 +78,8 @@ def convert_state_dict(
 
     ``state`` is either a raw state dict or a checkpoint holding it under the
     ``"model"`` key. Arrays keep their Torch layout and are returned C-contiguous.
+    Validation checks the parameter inventory and values, not the checkpoint's
+    identity or origin; other checkpoints with the same topology are admissible.
 
     Raises:
         InvalidInputError: If a required parameter is missing, an unsupported one is
@@ -117,11 +119,26 @@ def convert_checkpoint(checkpoint_path: Path, destination: Path, *, spec: SegRes
 
     The checkpoint is read with the shared loader (``torch.load`` with its default
     ``weights_only=True``, so Torch from the ``conversion`` extra is required) and
-    checked by ``convert_state_dict``. Parent directories of ``destination`` are
-    created. MLX is also needed to write the file, so this runs only on a supported
-    Apple Silicon host.
+    checked by ``convert_state_dict``. ``destination`` must end in
+    ``.safetensors``. Parent directories of ``destination`` are
+    created. An existing destination is replaced, except when it aliases the
+    source checkpoint. Alias checks occur before conversion and do not lock out
+    subsequent filesystem changes. MLX is also needed to write the file, so this
+    runs only on a supported Apple Silicon host.
+
+    Raises:
+        InvalidInputError: The destination has a different suffix or aliases the
+            source checkpoint, or checkpoint contents fail ``convert_state_dict``.
     """
 
+    if destination.suffix != ".safetensors":
+        raise InvalidInputError("Converted weights destination must end in .safetensors")
+    if checkpoint_path.resolve() == destination.resolve() or (
+        checkpoint_path.is_file()
+        and destination.is_file()
+        and checkpoint_path.samefile(destination)
+    ):
+        raise InvalidInputError("Converted weights must not overwrite their source checkpoint")
     mx = import_mlx()
     converted = convert_state_dict(load_torch_checkpoint(checkpoint_path), spec)
     destination.parent.mkdir(parents=True, exist_ok=True)
@@ -141,8 +158,9 @@ def load_model(
     """Load converted float32 safetensors into a ``SegResNet`` for ``spec``.
 
     Raises:
-        InvalidInputError: If the file is not a mapping, its keys or shapes differ
+        InvalidInputError: If the loaded value is not a mapping, its keys or shapes differ
             from ``required_shapes(spec)``, or a parameter is not finite float32.
+        RuntimeError: If MLX cannot open or decode the weights file.
     """
     mx = import_mlx()
     raw: object = mx.load(str(weights_path))