medmlx-core

Shared MLX building blocks for MedMLX medical imaging models on Apple Silicon.

• Metal runtime checks and memory reporting
• 3D convolutions, resampling and sliding-window inference
• Tensor layouts, precision helpers and checkpoint conversion
• Portable saved-graph inference and dense registration geometry

PyPI · Releases · Technical reference

Install · Python · Technical details · License

Requirements

Apple Silicon Mac, macOS, Python 3.13 or 3.14, and Metal GPU access. NumPy and MLX install automatically. Inference raises an error when Metal is unavailable.

Install

pip install medmlx-core

For optional PyTorch checkpoint conversion:

pip install "medmlx-core[conversion]"

Python

Run a small synthetic volume through overlapping patches:

import numpy as np
from medmlx_core import sliding_window_inference

volume = np.ones((1, 1, 32, 32, 32), dtype=np.float32)
output = sliding_window_inference(
    volume,
    roi_size=(16, 16, 16),
    sw_batch_size=1,
    predictor=lambda patch: patch * 2,
)
print(output.shape)  # (1, 1, 32, 32, 32)

Arrays use batch, channel, depth, height and width axes. Replace the example predictor with your model; the result is a NumPy array.

MedMLX model packages provide the architectures, weights and image workflows. This library supplies their shared runtime operations.

Converted graphs run directly through core:

from pathlib import Path
import mlx.core as mx
from medmlx_core.graph.runtime import Session

session = Session.load(Path("build/my-native-lungmask"))
log_probs, = session(mx.zeros((1, 1, 256, 256), dtype=mx.float32))

Graph schema 2, its codec and Metal executor are owned here. Source capture, ATen translation and source comparisons live in medmlx-convert. Inference does not require that converter, Torch, TensorFlow, Core ML or Core AI. Existing schema-1 pilot bundles must be regenerated; no legacy graph loader is retained. See the graph and geometry contract for the bounded operation coverage and numerical qualification scope.

medmlx_core.compensated provides explicit accuracy profiles using FP32 Metal storage. CompensatedFloat32Operators keeps product and summation residuals inside convolution, GroupNorm, frozen BatchNorm, ELU, softmax and Gaussian filtering, with FP32 outputs. CompensatedFloat48Operators carries both words between convolution, GroupNorm, transpose-convolution, addition and ReLU layers. The latter packs high and low words along the batch axis. Independent float64 source operators are the numerical references for both profiles.

For paired network outputs, sliding_window_inference(..., predictor_input="numpy") accepts a host NumPy callback and float64 windows and stitching. This explicit mode lets the callback run a paired Metal network and export both words for host blending. The default callback still uses MLX arrays and FP16/FP32 inputs.

Technical details

API contracts, numerical references and development commands are in the technical reference. See the decoder note for the Metal transpose-convolution fix.

License

Apache-2.0. Includes third-party notices and the upstream MONAI license.

Implementation mapping

implementation-map.json records upstream-to-native changes, source references and qualification scope.