| |
| |
| |
| |
| @@ -69,10 +69,11 @@ Inputs are 3D NIfTI volumes, all on the same voxel grid. Each run writes |
| ## Technical details |
| |
| Preprocessing, sliding-window inference and postprocessing follow the pinned MONAI |
| -1.6.0 bundles. Networks and runners are checked against seeded upstream recordings; |
| -saved masks match exactly. Real-checkpoint parity and clinical performance are not |
| -yet qualified. The lower-level API, tolerances and fixture regeneration are in |
| -[docs/technical.md](docs/technical.md). |
| +1.6.0 bundles. Networks and runners are checked against seeded upstream recordings. |
| +The released renal checkpoint also matches the independent MONAI graph exactly |
| +on the retained three-phase CT case, including raw logits and the saved mask. |
| +Clinical performance remains unqualified. The lower-level API, tolerances and |
| +fixture regeneration are in [docs/technical.md](docs/technical.md). |
| |
| ## Intended use |
| |
| |
| |
| |
| |
| @@ -19,17 +19,24 @@ MRI. Runtime code uses `medmlx-core`, NumPy and nibabel. Torch and MONAI are use |
| only for conversion or reference generation. Importing the package or its model |
| declaration does not import MLX. |
| |
| -**Qualification:** The recorded reference is the original MONAI implementation, |
| -with synthetic seeded inputs and weights. Renal normalized intensities are |
| -checked within one float32 rounding unit on [0, 1]; BraTS normalized intensities |
| -within `2**-20` (about 9.54e-7). All saved masks agree exactly. The reviewer's |
| -Metal run passed the seeded public-network, operator, sliding-window and runner |
| -logit checks at a 1e-6 absolute bound. Its sole failure was an encoder |
| -intermediate: 1.66893e-6 (seven float32 ULPs). Encoder intermediates now use a |
| -`2**-19` absolute bound (about 1.91e-6), reflecting native accumulation rounding; |
| -that changed assertion passes the current Metal suite (99 tests passed, two |
| -released-checkpoint tests skipped). Real-checkpoint numerical parity, |
| -clinical accuracy and full production-topology numerical parity remain unqualified. |
| +**Qualification:** The reference is the original MONAI implementation. Seeded |
| +public-network, operator, sliding-window and runner logit checks retain their |
| +1e-6 absolute bound; encoder intermediates retain `2**-19`. Renal normalized |
| +intensities are checked within one float32 rounding unit on [0, 1], and BraTS |
| +normalized intensities within `2**-20` (about 9.54e-7). Saved fixture masks agree |
| +exactly. |
| + |
| +The released renal checkpoint is also checked on dense and sparse seeded inputs |
| +with unchanged `rtol=atol=1e-4` and exact sign agreement. On 12 October 2026, |
| +source-order FP32 convolution and GroupNorm removed the previously retained |
| +sparse/full-case discrepancy. On the retained external three-phase CT case, |
| +all 85,089,840 raw logits were bitwise identical to a fresh independently parsed |
| +MONAI CPU bundle run, and all 5,935,600 exported mask voxels agreed. This was |
| +measured on Apple M1 Max, macOS 27.0.1, MLX 0.32.3, Torch 2.14.1 and MONAI 1.6.0; |
| +bitwise equality is an observation for those runs, not a promise for every |
| +platform or input. The case has no reference labels and retains manual |
| +registration/cropping and a denoising substitution. Anatomical and clinical |
| +accuracy remain unqualified. |
| |
| ## Install |
| |
| @@ -60,7 +67,11 @@ inference sections committed in `tests/fixtures/bundle_configs/` match those |
| original bundle files. The network uses their inference graph in float32 with |
| dropout disabled, group normalization, ReLU, skip addition, and either transposed |
| convolution or trilinear 2x interpolation (`align_corners=False`). Native MLX |
| -operations come from the shared core; no custom decoder kernels are required. |
| +operations come from the shared core. Convolution and GroupNorm use its |
| +source-order FP32 operators: channel-major fused multiply-add accumulation and |
| +bounded four-lane Welford moments. Correcting only one of these operators leaves |
| +amplified drift in the released renal checkpoint. The decoder uses the core's |
| +phase-packed transposed convolution or trilinear interpolation. |
| |
| The BraTS bundle enables mixed-precision evaluation upstream. The reference |
| contract uses its float32 evaluation graph; mixed-precision arithmetic is |
| @@ -160,7 +171,9 @@ qform/sform agreement rules or rewrite input files. |
| Renal preprocessing orders artery/vein/excretory channels, reorients to LPS, |
| resamples to 0.8 mm, and clips/scales [-1000, 1000] HU to [0, 1]. Postprocessing |
| inverts the spatial transforms, applies sigmoid and threshold `>= 0.5`, then |
| -merges classes with the lowest nonzero channel winning. Labels 1 artery, 2 vein, |
| +merges classes with the lowest nonzero channel winning. Sigmoid evaluates its |
| +negative tail without exponent overflow while preserving the source float32 |
| +expression near the 0.5 threshold. Labels 1 artery, 2 vein, |
| 3 ureter, 4 neoplasm and 5 parenchyma are saved as float32. |
| |
| BraTS normalizes nonzero voxels per modality, applies sigmoid, inversion and |
| @@ -224,4 +237,16 @@ writer, provenance, conversion-property and telemetry tests can run without it. |
| Use `-m 'not mlx'` with pytest to run that portion of the suite in a sandbox |
| without Metal. Full-suite collection does not load MLX; selecting a Metal test |
| loads the shared runtime when the test executes. |
| -There is no configured pyright check in this repository. |
| +The strict Pyright and annotation checks are configured in `make typecheck`. |
| +To include the released-checkpoint regression tests, stage the original bundle |
| +and run: |
| + |
| +```bash |
| +MEDMLX_TEST_CHECKPOINT_ROOT=/path/to/model-zoo \ |
| + uv run --locked --extra dev python -m pytest -q tests/test_torch_oracles.py |
| +``` |
| + |
| +The root must contain `renalStructures_CECT_segmentation/models/model.pt`. |
| +These checks use two independent seeds for both the dense 8-cube and sparse |
| +96-cube controls. Without this explicit asset root, those four tests are skipped; |
| +the remaining fixture and runtime tests still run. |
| |
| |
| |
| |
| @@ -9,7 +9,8 @@ if TYPE_CHECKING: |
| from ._types import MLX, Array |
| |
| from medmlx_core.layout import require_ncdhw |
| -from medmlx_core.ops import conv3d_ncdhw, conv_transpose3d_ncdhw, upsample_trilinear_ncdhw |
| +from medmlx_core.ops import conv_transpose3d_ncdhw, upsample_trilinear_ncdhw |
| +from medmlx_core.precision import Float32Operators |
| |
| from segresnet_mlx.topology import SegResNetSpec |
| |
| @@ -35,6 +36,9 @@ class SegResNet: |
| self.spec = spec |
| self._weights = weights |
| self._mx = mx |
| + # Both convolution and normalization accumulation order affect sparse |
| + # released-checkpoint logits; correcting only one leaves amplified drift. |
| + self._operators = Float32Operators(mx, spec.num_groups, spec.norm_eps) |
| |
| def __call__(self, values: Array) -> Array: |
| require_ncdhw(values, name="SegResNet input") |
| @@ -89,19 +93,9 @@ class SegResNet: |
| return hidden + values |
| |
| def _norm(self, values: Array, prefix: str) -> Array: |
| - mx = self._mx |
| - groups = self.spec.num_groups |
| - grouped = mx.reshape(values, (values.shape[0], groups, -1)) |
| - # A single long reduction loses precision on sparse full-volume features. |
| - mean = group_mean(grouped, mx) |
| - centered = grouped - mean |
| - variance = group_mean(centered * centered, mx) |
| - scale = mx.rsqrt(variance + self.spec.norm_eps) * mx.reshape( |
| - self._weights[f"{prefix}.weight"], (1, groups, -1) |
| + return self._operators.norm( |
| + values, self._weights[f"{prefix}.weight"], self._weights[f"{prefix}.bias"] |
| ) |
| - bias = mx.reshape(self._weights[f"{prefix}.bias"], (1, groups, -1)) - mean * scale |
| - shape = (values.shape[0], values.shape[1], 1, 1, 1) |
| - return values * mx.reshape(scale, shape) + mx.reshape(bias, shape) |
| |
| def conv( |
| self, |
| @@ -112,11 +106,10 @@ class SegResNet: |
| stride: int = 1, |
| bias: bool = False, |
| ) -> Array: |
| - return conv3d_ncdhw( |
| + return self._operators.conv( |
| values, |
| self._weights[f"{prefix}.weight"], |
| self._weights.get(f"{prefix}.bias") if bias else None, |
| padding=padding, |
| stride=stride, |
| - mx=self._mx, |
| ) |
| |
| |
| |
| |
| @@ -291,9 +291,16 @@ def normalize_intensity_nonzero(data: Float32Array) -> Float32Array: |
| |
| |
| def sigmoid(data: Float32Array) -> Float32Array: |
| - """``Activations(sigmoid=True)``: ``1 / (1 + exp(-x))`` in float32.""" |
| + """Float32 sigmoid with a stable negative tail and source threshold rounding.""" |
| one = np.float32(1.0) |
| - return (one / (one + np.exp(-data))).astype(np.float32, copy=False) |
| + result = np.empty_like(data) |
| + negative_tail = data < np.float32(-1.0) |
| + tail = np.exp(data[negative_tail]) |
| + result[negative_tail] = tail / (one + tail) |
| + # Keep the original expression near zero: equivalent sigmoid formulas can |
| + # round to opposite sides of the bundle's >= 0.5 threshold there. |
| + result[~negative_tail] = one / (one + np.exp(-data[~negative_tail])) |
| + return result |
| |
| |
| def threshold(data: Float32Array, value: float = 0.5) -> Float32Array: |
| |
| |
| |
| |
| @@ -25,8 +25,9 @@ _RAW_LOGIT_ATOL = 1e-4 |
| |
| |
| @pytest.mark.parametrize("sparse_volume", [False, True], ids=["dense", "sparse-accumulation"]) |
| +@pytest.mark.parametrize("seed", [23, 47]) |
| def test_renal_segresnet_checkpoint_matches_monai_raw_logits( |
| - tmp_path: Path, sparse_volume: bool |
| + tmp_path: Path, sparse_volume: bool, seed: int |
| ) -> None: |
| checkpoint_root = os.environ.get("MEDMLX_TEST_CHECKPOINT_ROOT") |
| if checkpoint_root is None: |
| @@ -60,7 +61,7 @@ def test_renal_segresnet_checkpoint_matches_monai_raw_logits( |
| state = state.get("model", state) |
| source.load_state_dict(state, strict=True) |
| |
| - rng = np.random.default_rng(23) |
| + rng = np.random.default_rng(seed) |
| if sparse_volume: |
| # Long sparse groups exposed reduction drift missed by the small controls. |
| size = 96 |
| @@ -84,6 +85,6 @@ def test_renal_segresnet_checkpoint_matches_monai_raw_logits( |
| assert host.shape == expected.shape |
| assert host.dtype == np.float32 |
| assert np.isfinite(host).all() |
| - assert np.allclose(host, expected, rtol=_RAW_LOGIT_RTOL, atol=_RAW_LOGIT_ATOL) |
| + np.testing.assert_allclose(host, expected, rtol=_RAW_LOGIT_RTOL, atol=_RAW_LOGIT_ATOL) |
| |
| assert np.array_equal(host >= 0, expected >= 0) |
| |
| |
| |
| |
| @@ -8,6 +8,7 @@ from typing import cast |
| |
| import numpy as np |
| import pytest |
| +from _contracts import from_numpy |
| from hypothesis import given, settings |
| from hypothesis import strategies as st |
| from reference_data import ( |
| @@ -18,7 +19,7 @@ from reference_data import ( |
| write_inputs, |
| ) |
| |
| -from segresnet_mlx import medmlx |
| +from segresnet_mlx import medmlx, transforms |
| from segresnet_mlx._arrays import MaskArray |
| from segresnet_mlx._nifti import nib |
| from segresnet_mlx.medmlx import BUNDLES |
| @@ -30,6 +31,21 @@ RENAL_KEYS = BUNDLE_INPUT_KEYS["renal"] |
| BRATS_KEYS = BUNDLE_INPUT_KEYS["brats"] |
| |
| |
| +def test_sigmoid_avoids_overflow_and_preserves_source_threshold_rounding() -> None: |
| + limit = np.finfo(np.float32).max |
| + values = np.array( |
| + [-limit, -1000, -100, -90, -80, -1, -1e-7, 0, 1e-7, 1, 80, 90, 100, limit], |
| + dtype=np.float32, |
| + ) |
| + expected = from_numpy(values).sigmoid().numpy() |
| + with np.errstate(over="raise", invalid="raise"): |
| + actual = transforms.sigmoid(values) |
| + assert actual.dtype == np.float32 |
| + assert np.isfinite(actual).all() |
| + np.testing.assert_allclose(actual, expected, rtol=0, atol=1e-6) |
| + np.testing.assert_array_equal(transforms.threshold(actual), expected >= 0.5) |
| + |
| + |
| @pytest.mark.parametrize("case", CASES) |
| def test_renal_transforms_match_the_bundle(tmp_path: Path, case: str) -> None: |
| arrays, _ = bundle_fixture("renal") |
|
|
