flexray-mlx 0.2.1__py3-none-any.whl

flexray_mlx/__init__.py

@@ -0,0 +1,21 @@
+"""Native MLX implementation of FleXray."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from flexray_mlx.mlx_model import FleXrayMlxModel
+
+__all__ = ["FleXrayMlxModel"]
+
+__version__ = "0.2.1"
+
+
+def __getattr__(name: str) -> object:
+    # MLX loads with the network, so importing the medmlx declaration stays light.
+    if name == "FleXrayMlxModel":
+        from flexray_mlx.mlx_model import FleXrayMlxModel
+
+        return FleXrayMlxModel
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

flexray_mlx/_kernels.py

@@ -0,0 +1,98 @@
+"""FP32 Metal fusions for the released FleXray decoder's NHWC tensors."""
+
+from __future__ import annotations
+
+from functools import lru_cache
+from typing import TYPE_CHECKING, cast
+
+from medmlx_core.typing import MetalKernel
+
+from flexray_mlx._types import NativeRuntime
+
+if TYPE_CHECKING:
+    from mlx.core import array
+
+
+_RESIZE_CONCAT = """
+// Preserve the original graph's separate multiply/add rounding.
+#pragma clang fp contract(off)
+uint i = thread_position_in_grid.x;
+constexpr int C = CI + CS;
+if (i >= N * OH * OW * C) return;
+uint c = i % C;
+uint ox = (i / C) % OW;
+uint oy = (i / (C * OW)) % OH;
+uint n = i / (C * OW * OH);
+if (c >= CI) {
+    out[i] = skip[((n * OH + oy) * OW + ox) * CS + c - CI];
+    return;
+}
+// align_corners=True: match the FP32 source-to-destination endpoint ratio.
+float iy = float(oy) * (float(H - 1) / float(OH - 1));
+float ix = float(ox) * (float(W - 1) / float(OW - 1));
+iy = metal::clamp(iy, 0.0f, float(H - 1));
+ix = metal::clamp(ix, 0.0f, float(W - 1));
+uint y0 = uint(metal::floor(iy)), y1 = uint(metal::ceil(iy));
+uint x0 = uint(metal::floor(ix)), x1 = uint(metal::ceil(ix));
+float fy = iy - float(y0), fx = ix - float(x0);
+float a = inp[((n * H + y0) * W + x0) * CI + c];
+float b = inp[((n * H + y0) * W + x1) * CI + c];
+float d = inp[((n * H + y1) * W + x0) * CI + c];
+float e = inp[((n * H + y1) * W + x1) * CI + c];
+out[i] = ((1.0f - fy) * (1.0f - fx)) * a
+       + ((1.0f - fy) * fx) * b
+       + (fy * (1.0f - fx)) * d
+       + (fy * fx) * e;
+"""
+
+
+@lru_cache(maxsize=1)
+def _resize_concat_kernel(mx: NativeRuntime) -> MetalKernel:
+    return cast(
+        MetalKernel,
+        mx.fast.metal_kernel(
+            name="flexray_bilinear_concat",
+            input_names=["inp", "skip"],
+            output_names=["out"],
+            source=_RESIZE_CONCAT,
+            compile_options={"math_mode": "safe"},
+        ),
+    )
+
+
+def upsample_concat(x: array, skip: array, *, mx: NativeRuntime) -> array:
+    """Join a 2x bilinear NHWC resize with ``skip``, keeping resized channels first.
+
+    Inputs must be FP32, with the same positive batch size and a skip grid
+    exactly twice ``x`` in both spatial dimensions. Corner-aligned sampling
+    matches the released decoder, including singleton source dimensions.
+    One kernel writes the joined tensor without materializing the four bilinear
+    sample tensors or the separate resized image. Noncontiguous inputs are
+    made contiguous by MLX's custom-kernel boundary.
+
+    Raises ``ValueError`` when the output exceeds signed 32-bit indexing, before
+    dispatch. The output is larger than either input, so this also bounds their
+    index products.
+    """
+    if x.ndim != 4 or skip.ndim != 4:
+        raise ValueError("FleXray decoder fusion requires NHWC tensors")
+    n, h, w, ci = x.shape
+    if min(n, h, w, ci, skip.shape[-1]) < 1 or skip.shape[:3] != (n, 2 * h, 2 * w):
+        raise ValueError("FleXray MLX decoder grid mismatch")
+    if x.dtype != mx.float32 or skip.dtype != mx.float32:
+        raise TypeError("FleXray decoder fusion requires float32 inputs")
+    oh, ow, cs = skip.shape[1:]
+    shape = (n, oh, ow, ci + cs)
+    count = n * oh * ow * (ci + cs)
+    if count > 2**31 - 1:
+        raise ValueError("FleXray decoder fusion exceeds signed 32-bit indexing")
+    return _resize_concat_kernel(mx)(
+        inputs=[x, skip],
+        template=[
+            ("N", n), ("H", h), ("W", w), ("CI", ci), ("CS", cs), ("OH", oh), ("OW", ow),
+        ],
+        grid=(count, 1, 1),
+        threadgroup=(256, 1, 1),
+        output_shapes=[shape],
+        output_dtypes=[mx.float32],
+    )[0]

flexray_mlx/_random.py

@@ -0,0 +1,91 @@
+"""Torch CPU draw order for the pinned FleXray float32 augmentation chain.
+
+MT19937 uses its low 24 bits for float uniforms. Contiguous normal tensors use
+16-value Box-Muller blocks; short tensors use cached double-precision pairs.
+Transcendental rounding can differ between CPU libraries, as it does in Torch.
+Adapted from PyTorch's ATen/core and native/cpu distribution helpers (BSD;
+see root THIRD_PARTY_NOTICES.md).
+"""
+
+from __future__ import annotations
+
+import math
+
+import numpy as np
+
+from flexray_mlx._types import FloatArray, UIntArray
+
+
+class TorchRandom:
+    """An isolated stream that reproduces Torch CPU draws, without importing Torch.
+
+    Draws follow Torch's order of consumption for a given seed and call sequence;
+    the module docstring notes where float rounding can still differ. Values are
+    float32.
+    """
+
+    def __init__(self, seed: int | None) -> None:
+        """Seed from the low 32 bits of ``seed`` (negative values wrap); None is unseeded."""
+
+        self._engine = np.random.RandomState(None if seed is None else seed & 0xFFFFFFFF)
+        self._next_normal: float | None = None
+
+    def _words(self, count: int) -> UIntArray:
+        return self._engine.randint(0, 2**32, size=count, dtype=np.uint32)
+
+    def _units(self, count: int) -> FloatArray:
+        return (self._words(count) & np.uint32(0xFFFFFF)).astype(np.float32) * np.float32(2**-24)
+
+    def uniform(self, low: float = 0.0, high: float = 1.0) -> float:
+        """Draw one float32 uniform between ``low`` and ``high``, like Torch's ``uniform_``."""
+        # Deferred: flexray_mlx.medmlx imports this module (via tta) at load time.
+        from flexray_mlx.medmlx import fma32
+
+        lower, upper = np.float32(low), np.float32(high)
+        return float(fma32(self._units(1)[0], upper - lower, lower))
+
+    def chance(self, probability: float) -> np.bool_:
+        """Consume one uniform draw and return whether it is below ``probability``."""
+        return self.uniform() < np.float32(probability)
+
+    def _double_unit(self) -> float:
+        words = self._words(2).astype(np.uint64)
+        value = (int(words[0]) << 32) | int(words[1])
+        return (value & (2**53 - 1)) * 2**-53
+
+    def _short_normal(self) -> float:
+        if self._next_normal is not None:
+            value, self._next_normal = self._next_normal, None
+            return value
+        theta = 2.0 * math.pi * self._double_unit()
+        radius = math.sqrt(-2.0 * math.log1p(-self._double_unit()))
+        self._next_normal = radius * math.sin(theta)
+        return radius * math.cos(theta)
+
+    @staticmethod
+    def _normal_blocks(uniforms: FloatArray) -> FloatArray:
+        blocks = uniforms.reshape(-1, 16)
+        radius = np.sqrt(np.float32(-2.0) * np.log(np.float32(1.0) - blocks[:, :8]))
+        # The scalar CPU kernel multiplies by double pi before rounding theta.
+        theta = (2.0 * math.pi * blocks[:, 8:].astype(np.float64)).astype(np.float32)
+        return np.concatenate((radius * np.cos(theta), radius * np.sin(theta)), axis=1).ravel()
+
+    def normal(self, shape: tuple[int, ...]) -> FloatArray:
+        """Draw a float32 standard-normal array of ``shape``.
+
+        Fewer than 16 values use cached double-precision Box-Muller pairs; 16 or
+        more use float32 blocks of 16, as Torch's CPU kernel does.
+        """
+        count = math.prod(shape)
+        if count < 16:
+            return np.asarray([self._short_normal() for _ in range(count)], np.float32).reshape(
+                shape
+            )
+        uniforms = self._units(count)
+        complete = count // 16 * 16
+        values = np.empty(count, dtype=np.float32)
+        values[:complete] = self._normal_blocks(uniforms[:complete])
+        if count % 16:
+            # Torch redraws the final 16, including the overlapping completed block.
+            values[-16:] = self._normal_blocks(self._units(16))
+        return values.reshape(shape)

flexray_mlx/_types.py

@@ -0,0 +1,126 @@
+"""Array and namespace contracts for the pinned native and file boundaries."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from os import PathLike
+from typing import TYPE_CHECKING, Literal, Protocol, TypedDict
+
+import numpy as np
+from medmlx_core.typing import FastRuntime
+from numpy.typing import NDArray
+
+if TYPE_CHECKING:
+    from mlx.core import Dtype, array
+
+type HostArray = NDArray[np.generic]
+type RealArray = NDArray[np.bool_ | np.integer | np.floating]
+type FloatArray = NDArray[np.float32]
+type DoubleArray = NDArray[np.float64]
+type BoolArray = NDArray[np.bool_]
+type IntArray = NDArray[np.int64]
+type UIntArray = NDArray[np.uint32]
+type JsonValue = bool | int | float | str | list[JsonValue] | dict[str, JsonValue] | None
+type JsonObject = dict[str, JsonValue]
+
+
+class NativeRuntime(Protocol):
+    fast: FastRuntime
+    float32: Dtype
+
+    # Constructor inputs here are the NumPy arrays prepared by this network.
+    def array(self, value: FloatArray) -> array: ...
+    def eval(self, *args: array) -> None: ...
+    def conv2d(self, input: array, weight: array, /, *, padding: int) -> array: ...
+    def mean(self, a: array, /, *, axis: Sequence[int], keepdims: bool) -> array: ...
+    def square(self, a: array, /) -> array: ...
+    def rsqrt(self, a: array, /) -> array: ...
+    def maximum(self, a: array, b: array, /) -> array: ...
+    def concatenate(self, arrays: list[array], *, axis: int) -> array: ...
+
+
+class NativeOperation(Protocol):
+    def __call__(self, x: array) -> array: ...
+
+
+class NeuralRuntime(Protocol):
+    def MaxPool2d(self, *, kernel_size: int, stride: int) -> NativeOperation: ...
+    def Upsample(self, *, scale_factor: int, mode: str, align_corners: bool) -> NativeOperation: ...
+
+
+class TensorReader(Protocol):
+    def __enter__(self) -> TensorReader: ...
+    def __exit__(self, *args: object) -> None: ...
+    def keys(self) -> list[str]: ...
+    def get_tensor(self, name: str) -> HostArray: ...
+
+
+class SafetensorsRuntime(Protocol):
+    def safe_open(self, filename: PathLike[str], *, framework: Literal["np"]) -> TensorReader: ...
+
+
+class DicomDataset(Protocol):
+    pixel_array: RealArray
+
+    def __contains__(self, name: str) -> bool: ...
+
+
+class DicomRuntime(Protocol):
+    def dcmread(self, fp: str, *, force: bool) -> DicomDataset: ...
+
+
+class Normalization(TypedDict):
+    scheme: str
+    percentiles: tuple[float, float] | None
+    eps: float
+    mean: float | None
+    std: float | None
+
+
+class EnsembleMember(TypedDict):
+    subfolder: str
+
+
+class EnsembleManifest(TypedDict):
+    members: list[EnsembleMember]
+    flagship: str
+
+
+class RunOutputs(TypedDict):
+    """POSIX paths of the files ``run`` writes to its output directory."""
+
+    mask: str
+    display_source: str
+    probabilities: str
+    prediction: str
+
+
+class RunDetails(TypedDict):
+    """Settings and geometry of one ``run``.
+
+    ``source_size`` is ``[width, height]`` of the input; ``model_grid`` is
+    ``[height, width]``. ``content_box_tlbr`` is ``[top, left, bottom, right]``
+    in model-grid pixels with exclusive bottom/right.
+    """
+
+    selected_label: str
+    selected_channel: int
+    threshold: float
+    ensemble: bool
+    members: list[str]
+    tta_samples: int
+    tta_seed: int | None
+    normalization_scheme: str
+    source_size: list[int]
+    model_grid: list[int]
+    content_box_tlbr: list[int]
+    mask_pixels: int
+    probability_mode: str
+    backend: str
+
+
+class RunResult(TypedDict):
+    """Return value of ``run``: output file paths and run details."""
+
+    outputs: RunOutputs
+    details: RunDetails