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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: segresnet-mlx
-Version: 0.2.1
+Version: 0.2.2
 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,19 +47,12 @@ 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
@@ -80,12 +73,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.2}/README.md

@@ -18,19 +18,12 @@ 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
@@ -51,12 +44,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.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "segresnet-mlx"
-version = "0.2.1"
+version = "0.2.2"
 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.2}/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.2"
 
 _EXPORTS = {
     "BUNDLE_CONFIGS": "segresnet_mlx.predictor",

{segresnet_mlx-0.2.1 → segresnet_mlx-0.2.2}/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.2}/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
@@ -71,13 +71,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 +188,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 +250,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.2}/src/segresnet_mlx/transforms.py

@@ -517,12 +517,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 +534,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 +585,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 +612,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 +622,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 +641,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