lesiondiffusion-mlx 0.2.3__tar.gz → 0.2.4__tar.gz

{lesiondiffusion_mlx-0.2.3 → lesiondiffusion_mlx-0.2.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: lesiondiffusion-mlx
-Version: 0.2.3
+Version: 0.2.4
 Summary: Native MLX implementation of LesionDiffusion
 Project-URL: Homepage, https://huggingface.co/MedMLX/lesiondiffusion-mlx
 Author: Joseph Sandoval, Alfredo Sandoval Jr.
@@ -47,7 +47,7 @@ Add synthetic lesions to CT with native MLX, with controls for organ, size, plac
 ```bash
 uv venv --python 3.13
 source .venv/bin/activate
-uv pip install "lesiondiffusion-mlx>=0.2.3"
+uv pip install "lesiondiffusion-mlx>=0.2.4"
 ```
 
 ```python

{lesiondiffusion_mlx-0.2.3 → lesiondiffusion_mlx-0.2.4}/README.md

@@ -15,7 +15,7 @@ Add synthetic lesions to CT with native MLX, with controls for organ, size, plac
 ```bash
 uv venv --python 3.13
 source .venv/bin/activate
-uv pip install "lesiondiffusion-mlx>=0.2.3"
+uv pip install "lesiondiffusion-mlx>=0.2.4"
 ```
 
 ```python

{lesiondiffusion_mlx-0.2.3 → lesiondiffusion_mlx-0.2.4}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "lesiondiffusion-mlx"
-version = "0.2.3"
+version = "0.2.4"
 description = "Native MLX implementation of LesionDiffusion"
 readme = "README.md"
 requires-python = ">=3.13,<3.15"

{lesiondiffusion_mlx-0.2.3 → lesiondiffusion_mlx-0.2.4}/src/lesiondiffusion_mlx/__init__.py

@@ -32,7 +32,7 @@ if TYPE_CHECKING:
     )
     from lesiondiffusion_mlx.mlx_sampling import DdimSampler as DdimSampler
 
-__version__ = "0.2.3"
+__version__ = "0.2.4"
 
 __all__ = [
     "LINET_CONFIG",
@@ -48,6 +48,7 @@ __all__ = [
 
 
 def __getattr__(name: str) -> object:
+    """Load a public export from its module on first access."""
     if name not in __all__:
         raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
     module = (

{lesiondiffusion_mlx-0.2.3 → lesiondiffusion_mlx-0.2.4}/src/lesiondiffusion_mlx/_group_norm.py

@@ -98,6 +98,7 @@ if (SILU) y = y / (1.0f + metal::precise::exp(-y));
 
 @lru_cache(maxsize=1)
 def _kernels(mx: Mlx) -> tuple[MetalKernel, MetalKernel, MetalKernel]:
+    """Build or reuse the native group-normalization Metal kernels."""
     return (
         mx.fast.metal_kernel(
             name="lesion_group_norm_partial",

{lesiondiffusion_mlx-0.2.3 → lesiondiffusion_mlx-0.2.4}/src/lesiondiffusion_mlx/_nax_conv.py

@@ -392,6 +392,7 @@ def pack_upsample_weight(weight: Array, mx: Mlx) -> Array:
     # parity p reads padded source offset p + t. Plain FP32 additions, not an
     # einsum, whose matmul could run in TF32 under MLX_ENABLE_TF32.
     def fold(taps: Array) -> Array:
+        """Fold one upsampling tap into packed convolution weights."""
         k0, k1, k2 = taps[0], taps[1], taps[2]
         return mx.stack([mx.stack([k0, k1 + k2]), mx.stack([k0 + k1, k2])])
 

lesiondiffusion_mlx-0.2.4/src/lesiondiffusion_mlx/_nifti.py

@@ -0,0 +1,111 @@
+"""The NIfTI surface used here, verified against NiBabel's image/header APIs."""
+
+from os import PathLike
+from typing import Literal, Protocol, cast
+
+import nibabel
+import numpy as np
+from nibabel.filebasedimages import ImageFileError as ImageFileError
+
+from lesiondiffusion_mlx._arrays import HostArray
+
+type PathArg = str | PathLike[str]
+type SpatialAffine = np.ndarray[tuple[int, ...], np.dtype[np.float64]]
+
+
+class ArrayBuffer(Protocol):
+    def __array__(
+        self, dtype: np.dtype[np.generic] | None = None, copy: bool | None = None
+    ) -> HostArray:
+        """Convert the image data buffer to a host array."""
+        ...
+
+
+class NiftiHeader(Protocol):
+    def get_zooms(self) -> tuple[float, ...]:
+        """Return voxel spacing from the header."""
+        ...
+    def get_xyzt_units(self) -> tuple[str, str]:
+        """Return spatial and temporal unit names."""
+        ...
+    def set_xyzt_units(self, xyz: str, t: str) -> None:
+        """Set spatial and temporal unit names."""
+        ...
+    def set_data_dtype(self, dtype: type[np.generic]) -> None:
+        """Set the stored image dtype."""
+        ...
+
+
+class HeaderFactory(Protocol):
+    def from_header(self, header: NiftiHeader) -> NiftiHeader:
+        """Construct a NIfTI header from an existing header."""
+        ...
+
+
+class NiftiImage(Protocol):
+    shape: tuple[int, ...]
+    affine: SpatialAffine
+    dataobj: ArrayBuffer | HostArray
+    header: NiftiHeader
+
+    def get_qform(self, *, coded: Literal[True]) -> tuple[SpatialAffine | None, int]:
+        """Return the qform affine and its code."""
+        ...
+    def get_sform(self, *, coded: Literal[True]) -> tuple[SpatialAffine | None, int]:
+        """Return the sform affine and its code."""
+        ...
+    def set_qform(self, affine: HostArray, code: int) -> None:
+        """Set the qform affine and code."""
+        ...
+    def set_sform(self, affine: HostArray, code: int) -> None:
+        """Set the sform affine and code."""
+        ...
+
+
+class ImageFactory(Protocol):
+    def __call__(
+        self, dataobj: HostArray, affine: HostArray, header: NiftiHeader | None = None
+    ) -> NiftiImage:
+        """Construct a NIfTI image from data, affine, and header."""
+        ...
+
+
+class OrientationFunctions(Protocol):
+    """The subset of ``nibabel.orientations`` used to resample to RAS 1 mm."""
+
+    def axcodes2ornt(self, axcodes: tuple[str, str, str]) -> HostArray:
+        """Convert axis codes to an orientation transform."""
+        ...
+    def ornt_transform(self, start_ornt: HostArray, end_ornt: HostArray) -> HostArray:
+        """Compute the transform between two orientations."""
+        ...
+    def apply_orientation(self, data: HostArray, ornt: HostArray) -> HostArray:
+        """Apply an orientation transform to image data."""
+        ...
+    def inv_ornt_aff(self, ornt: HostArray, shape: tuple[int, ...]) -> SpatialAffine:
+        """Return the inverse voxel affine for an orientation and shape."""
+        ...
+
+
+class Nibabel(Protocol):
+    Nifti1Header: HeaderFactory
+    Nifti1Image: ImageFactory
+    orientations: OrientationFunctions
+
+    def load(self, filename: PathArg) -> NiftiImage:
+        """Load an image from a path."""
+        ...
+    def save(self, image: NiftiImage, filename: PathArg) -> None:
+        """Write a NIfTI image to a path."""
+        ...
+    def aff2axcodes(self, affine: HostArray) -> tuple[str | None, ...]:
+        """Derive axis codes from an affine."""
+        ...
+    def io_orientation(self, affine: HostArray) -> HostArray:
+        """Derive voxel-axis orientation from an affine."""
+        ...
+
+
+# This adapter is used only for .nii/.nii.gz inputs and NIfTI output; the wider
+# NiBabel loader supports unrelated image classes outside the runner contract.
+nib = cast(Nibabel, nibabel)

{lesiondiffusion_mlx-0.2.3 → lesiondiffusion_mlx-0.2.4}/src/lesiondiffusion_mlx/_reference.py

@@ -25,6 +25,7 @@ class SourceReference(TypedDict):
 
 
 def _string(value: object) -> str:
+    """Require a string for a reference metadata field."""
     if not isinstance(value, str):
         raise ValueError("reference.json source names, revisions and labels must be strings")
     return value

{lesiondiffusion_mlx-0.2.3 → lesiondiffusion_mlx-0.2.4}/src/lesiondiffusion_mlx/_schedules.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 
 def _is_runtime_integer(value: object) -> bool:
+    """Recognize integer inputs while excluding booleans."""
     return isinstance(value, int) and not isinstance(value, bool)
 
 

{lesiondiffusion_mlx-0.2.3 → lesiondiffusion_mlx-0.2.4}/src/lesiondiffusion_mlx/_source_half.py

@@ -342,9 +342,11 @@ def _norm_affine_kernels(mx: Mlx) -> tuple[MetalKernel, MetalKernel]:
 
 @lru_cache(maxsize=1)
 def _kernels(mx: Mlx) -> tuple[MetalKernel, MetalKernel, MetalKernel, MetalKernel, MetalKernel]:
+    """Build or reuse the source-order half-precision Metal kernels."""
     def kernel(
         name: str, inputs: list[str], output: str, source: str, header: str = ""
     ) -> MetalKernel:
+        """Construct one named source-order Metal kernel."""
         return mx.fast.metal_kernel(
             name="lesion_source_half_" + name,
             input_names=inputs,

lesiondiffusion_mlx-0.2.4/src/lesiondiffusion_mlx/_typing.py

@@ -0,0 +1,178 @@
+"""Native array and dependency contracts used by the source implementation."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Literal, Protocol
+
+from tokenizers import Encoding
+
+from lesiondiffusion_mlx._arrays import HostArray as HostArray
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Sequence
+
+    import mlx.core as mx
+    from medmlx_core.typing import FastRuntime as CoreFastRuntime
+    from medmlx_core.typing import MlxRuntime
+
+    type Array = mx.array
+    type Dtype = mx.Dtype
+    type ArrayValue = bool | int | float | complex | list[ArrayValue] | tuple[ArrayValue, ...]
+
+    class ArrayConstructor(Protocol):
+        def __call__(
+            self,
+            value: HostArray | Array | ArrayValue,
+            /,
+            dtype: Dtype | None = None,
+        ) -> Array:
+            """Construct a native array from values and an optional dtype."""
+            ...
+
+    class TokenizerContract(Protocol):
+        def enable_truncation(self, max_length: int) -> None:
+            """Enable tokenizer truncation at the requested length."""
+            ...
+        def enable_padding(self, *, length: int, pad_id: int, pad_token: str) -> None:
+            """Configure tokenizer padding and its token ID."""
+            ...
+        def encode_batch(self, inputs: Sequence[str]) -> list[Encoding]:
+            """Encode a batch of input strings."""
+            ...
+
+    class MetalKernel(Protocol):
+        def __call__(
+            self,
+            *,
+            inputs: Sequence[Array],
+            template: Sequence[tuple[str, int | bool | Dtype]],
+            grid: tuple[int, int, int],
+            threadgroup: tuple[int, int, int],
+            output_shapes: Sequence[Sequence[int]],
+            output_dtypes: Sequence[Dtype],
+        ) -> list[Array]:
+            """Launch a Metal kernel and return its output arrays."""
+            ...
+
+    class Fast(CoreFastRuntime, Protocol):
+        def scaled_dot_product_attention(
+            self,
+            q: Array,
+            k: Array,
+            v: Array,
+            *,
+            scale: float,
+            mask: str | Array | None = None,
+            sinks: Array | None = None,
+            force_fused: bool = False,
+        ) -> Array:
+            """Compute native scaled dot-product attention."""
+            ...
+
+        def metal_kernel(
+            self,
+            name: str,
+            input_names: Sequence[str],
+            output_names: Sequence[str],
+            source: str,
+            header: str = "",
+            ensure_row_contiguous: bool = True,
+            atomic_outputs: bool = False,
+            compile_options: object | None = None,
+        ) -> MetalKernel:
+            """Compile a named native Metal kernel."""
+            ...
+
+    class Random(Protocol):
+        def seed(self, seed: int) -> None:
+            """Set the native random generator seed."""
+            ...
+        def normal(self, shape: Sequence[int]) -> Array:
+            """Draw normally distributed native values."""
+            ...
+
+    class Mlx(MlxRuntime, Protocol):
+        """The MLX module surface consumed by this port, without importing it eagerly."""
+
+        @property
+        def fast(self) -> Fast:
+            """Return the native fast operations boundary."""
+            ...
+
+        random: Random
+
+        def addmm(self, c: Array, a: Array, b: Array, /) -> Array:
+            """Multiply matrices and add the supplied bias."""
+            ...
+        def all(self, a: Array, /) -> Array:
+            """Reduce a native array with logical all."""
+            ...
+        def argmin(self, a: Array, /, *, axis: int | None = None) -> Array:
+            """Return indices of minimum values along an axis."""
+            ...
+        def clear_cache(self) -> None:
+            """Clear the native allocator cache."""
+            ...
+        def compile[**P, R](self, fun: Callable[P, R], /) -> Callable[P, R]:
+            """Compile a callable for native execution."""
+            ...
+        def contiguous(self, a: Array, /) -> Array:
+            """Return a contiguous native array."""
+            ...
+        def cos(self, a: Array, /) -> Array:
+            """Compute elementwise cosine."""
+            ...
+        def device_info(self) -> dict[str, str | int]:
+            """Return native device capability information."""
+            ...
+        def erf(self, a: Array, /) -> Array:
+            """Compute the elementwise error function."""
+            ...
+        def exp(self, a: Array, /) -> Array:
+            """Compute the elementwise exponential."""
+            ...
+        def full(self, shape: Sequence[int], vals: object, *, dtype: Dtype) -> Array:
+            """Create a native array filled with the supplied value."""
+            ...
+        def isfinite(self, a: Array, /) -> Array:
+            """Test native values for finiteness."""
+            ...
+        def load(self, file: str, *, format: Literal["safetensors"]) -> dict[str, Array]:
+            """Load named arrays from a safetensors file."""
+            ...
+        def reciprocal(self, a: Array, /) -> Array:
+            """Compute elementwise reciprocals."""
+            ...
+        def sin(self, a: Array, /) -> Array:
+            """Compute elementwise sine."""
+            ...
+        def softmax(
+            self, a: Array, /, *, axis: int | None = None, precise: bool = False
+        ) -> Array:
+            """Normalize scores along the requested axis."""
+            ...
+        def sqrt(self, a: Array | float, /) -> Array:
+            """Compute elementwise square roots."""
+            ...
+        def sum(self, a: Array, /, *, axis: int | None = None, keepdims: bool = False) -> Array:
+            """Sum native values over an optional axis."""
+            ...
+else:
+
+    class ArrayConstructor(Protocol):
+        pass
+
+    class TokenizerContract(Protocol):
+        pass
+
+    class MetalKernel(Protocol):
+        pass
+
+    class Fast(Protocol):
+        pass
+
+    class Random(Protocol):
+        pass
+
+    class Mlx(Protocol):
+        pass

{lesiondiffusion_mlx-0.2.3 → lesiondiffusion_mlx-0.2.4}/src/lesiondiffusion_mlx/medmlx.py

@@ -249,6 +249,7 @@ class LesionDiffusionOptions:
     output_region: str = DEFAULT_OUTPUT_REGION
 
     def __post_init__(self) -> None:
+        """Validate inference, geometry, and lesion-generation options."""
         if self.inference_mode not in {"conditioned", "source"}:
             raise InvalidInputError("LesionDiffusion inference_mode must be conditioned or source")
         if self.output_region not in {"full", "box"}:
@@ -495,6 +496,7 @@ def _run_paths(
     output_region: str,
     overwrite: bool,
 ) -> tuple[Path, Path, dict[str, Path]]:
+    """Resolve input and output paths and prevent output collisions."""
     if not isinstance(cast(object, overwrite), bool):
         raise InvalidInputError("overwrite must be a boolean")
     source = Path(image_path).expanduser().resolve()
@@ -542,6 +544,7 @@ def _execute_run(
     directory: Path,
     outputs: dict[str, Path],
 ) -> tuple[GenerationSpec, dict[str, dict[str, object]]]:
+    """Run native generation and write its validated outputs."""
     model_label, _ = source_organ_label(options.organ)
     report = _structured_report(
         reference=geometry.reference,
@@ -603,6 +606,7 @@ def _run_result(
     started_at: datetime,
     started: float,
 ) -> RunResult:
+    """Build the output paths and generation metadata for a completed run."""
     grid = [int(value) for value in geometry.output_reference.shape]
     details: RunDetails = {
         "model_id": MODEL_ID,
@@ -767,6 +771,7 @@ def run_native_pipeline(
 def _linet_attributes(
     attributes: Mapping[str, object], options: LesionDiffusionOptions
 ) -> dict[str, object]:
+    """Build image-synthesis text attributes for the selected organ and mode."""
     _, organ_name = source_organ_label(options.organ)
     omitted = {"size"}
     if options.inference_mode == "conditioned":
@@ -787,6 +792,7 @@ def _lmnet_attributes(
     attributes: Mapping[str, object],
     options: LesionDiffusionOptions,
 ) -> dict[str, object]:
+    """Build mask-synthesis text attributes using the pinned source conventions."""
     _, organ_name = source_organ_label(options.organ)
     lm_attributes = dict(attributes)
     lm_attributes["lesion location"] = organ_name
@@ -822,6 +828,7 @@ def _run_lmnet(
     *,
     anatomy: HostArray,
 ) -> tuple[BoolArray, dict[str, object]]:
+    """Generate and validate a lesion mask within the selected crop."""
     crop = source_crop(bbox, LMNET_CROP)
     lm_started = time.perf_counter()
     optimized = options.inference_mode == "conditioned"
@@ -890,6 +897,7 @@ def _run_linet(
     tokenizer_path: Path,
     mx: Mlx,
 ) -> tuple[HostArray, dict[str, object]]:
+    """Generate the lesion image using the selected mask and attributes."""
     lin_started = time.perf_counter()
     source_arithmetic = options.inference_mode == "source"
     linet = LdmUnet(
@@ -1156,6 +1164,7 @@ def load_stage(directory: Path, name: str, mx: Mlx, *, dtype: Dtype) -> dict[str
 
 
 def _load_json_object(path: Path) -> dict[str, object]:
+    """Read a JSON object or report an invalid metadata file."""
     try:
         payload: object = json.loads(path.read_text(encoding="utf-8"))
     except (UnicodeError, json.JSONDecodeError) as exc:
@@ -1188,6 +1197,7 @@ def _parse_options(
     output_region: object = DEFAULT_OUTPUT_REGION,
     lesion_mask_supplied: bool = False,
 ) -> LesionDiffusionOptions:
+    """Validate request values and construct generation options."""
     for name, value in (
         ("organ", organ),
         ("lesion_type", lesion_type),
@@ -1274,6 +1284,7 @@ def _check_box_source(
 
 
 def _is_integer(value: object) -> bool:
+    """Recognize integer inputs while excluding booleans."""
     return isinstance(value, int) and not isinstance(value, bool)
 
 
@@ -2137,11 +2148,13 @@ def localize_output(
 
 def _with_comma(value: str) -> str:
     # Upstream report phrases are comma-terminated.
+    """Normalize an upstream report phrase to end with one comma."""
     stripped = value.strip().rstrip(",")
     return f"{stripped},"
 
 
 def _json_safe(value: object) -> object:
+    """Convert nested mappings, sequences, and NumPy scalars to JSON-compatible values."""
     if isinstance(value, Mapping):
         return {
             str(key): _json_safe(item) for key, item in cast(Mapping[object, object], value).items()

{lesiondiffusion_mlx-0.2.3 → lesiondiffusion_mlx-0.2.4}/src/lesiondiffusion_mlx/mlx_model.py

@@ -99,6 +99,7 @@ class SourceLayers:
         source_arithmetic: bool = True,
         native_kernels: bool | None = None,
     ):
+        """Initialize source-keyed weights and the selected native arithmetic."""
         self.weights = weights
         self.mx = mx
         self.source_arithmetic = source_arithmetic
@@ -188,6 +189,7 @@ class SourceLayers:
         return linear(x, weight, bias, self.mx).astype(dtype)
 
     def _affine_inputs(self, x: Array, key: str) -> tuple[Array, Array, Array | None]:
+        """Promote half inputs and weights for source-compatible affine accumulation."""
         weight = self.weights[key + ".weight"]
         bias = self.weights.get(key + ".bias")
         if x.dtype == self.mx.float16:
@@ -313,6 +315,7 @@ class SourceLayers:
         return silu(x, self.mx)
 
     def _native_norm(self, x: Array, key: str, *, groups: int, eps: float, silu: bool) -> Array:
+        """Apply the native group-normalization kernel with optional SiLU."""
         return native_group_norm(
             x,
             self.weights[key + ".weight"],
@@ -374,6 +377,7 @@ class LdmUnet(SourceLayers):
         fused_attention: bool = False,
         native_kernels: bool | None = None,
     ):
+        """Initialize the source UNet and its attention and arithmetic policies."""
         super().__init__(
             weights, mx, source_arithmetic=source_arithmetic, native_kernels=native_kernels
         )
@@ -383,10 +387,12 @@ class LdmUnet(SourceLayers):
     def silu(self, x: Array) -> Array:
         # Source nn.SiLU rounds once after the activation. VQ's separate
         # sigmoid/product deliberately retains its two half rounding steps.
+        """Apply SiLU with FP32 intermediate rounding and restore the input dtype."""
         return silu(x.astype(self.mx.float32), self.mx).astype(x.dtype)
 
     def softmax(self, x: Array) -> Array:
         # Torch softmax accumulates in float32, then returns the input dtype.
+        """Apply source-compatible precise softmax on the final axis."""
         return self.mx.softmax(x, axis=-1, precise=True)
 
     def residual(self, x: Array, emb: Array, key: str) -> Array:
@@ -432,6 +438,7 @@ class LdmUnet(SourceLayers):
             return self.dense(result, key + ".to_out.0")
 
         def split(tensor: Array) -> Array:
+            """Split projected tokens into attention heads."""
             return mx.transpose(
                 mx.reshape(tensor, (tensor.shape[0], tensor.shape[1], heads, depth)), (0, 2, 1, 3)
             )
@@ -548,6 +555,7 @@ class VqAutoencoder(SourceLayers):
     """
 
     def residual(self, x: Array, key: str) -> Array:
+        """Apply one VQ autoencoder residual block."""
         h = self.norm_silu_conv(x, key + ".norm1", key + ".conv1", groups=16, eps=1e-6)
         if key + ".nin_shortcut.weight" in self.weights:
             x = self.conv(x, key + ".nin_shortcut", padding=0)
@@ -653,6 +661,7 @@ class BiomedTextEncoder(SourceLayers):
             )
 
             def split(tensor: Array) -> Array:
+                """Split text-attention projections into heads."""
                 return mx.transpose(
                     mx.reshape(tensor, (tensor.shape[0], tensor.shape[1], 12, 64)), (0, 2, 1, 3)
                 )

lesiondiffusion_mlx-0.2.3/src/lesiondiffusion_mlx/_nifti.py

@@ -1,73 +0,0 @@
-"""The NIfTI surface used here, verified against NiBabel's image/header APIs."""
-
-from os import PathLike
-from typing import Literal, Protocol, cast
-
-import nibabel
-import numpy as np
-from nibabel.filebasedimages import ImageFileError as ImageFileError
-
-from lesiondiffusion_mlx._arrays import HostArray
-
-type PathArg = str | PathLike[str]
-type SpatialAffine = np.ndarray[tuple[int, ...], np.dtype[np.float64]]
-
-
-class ArrayBuffer(Protocol):
-    def __array__(
-        self, dtype: np.dtype[np.generic] | None = None, copy: bool | None = None
-    ) -> HostArray: ...
-
-
-class NiftiHeader(Protocol):
-    def get_zooms(self) -> tuple[float, ...]: ...
-    def get_xyzt_units(self) -> tuple[str, str]: ...
-    def set_xyzt_units(self, xyz: str, t: str) -> None: ...
-    def set_data_dtype(self, dtype: type[np.generic]) -> None: ...
-
-
-class HeaderFactory(Protocol):
-    def from_header(self, header: NiftiHeader) -> NiftiHeader: ...
-
-
-class NiftiImage(Protocol):
-    shape: tuple[int, ...]
-    affine: SpatialAffine
-    dataobj: ArrayBuffer | HostArray
-    header: NiftiHeader
-
-    def get_qform(self, *, coded: Literal[True]) -> tuple[SpatialAffine | None, int]: ...
-    def get_sform(self, *, coded: Literal[True]) -> tuple[SpatialAffine | None, int]: ...
-    def set_qform(self, affine: HostArray, code: int) -> None: ...
-    def set_sform(self, affine: HostArray, code: int) -> None: ...
-
-
-class ImageFactory(Protocol):
-    def __call__(
-        self, dataobj: HostArray, affine: HostArray, header: NiftiHeader | None = None
-    ) -> NiftiImage: ...
-
-
-class OrientationFunctions(Protocol):
-    """The subset of ``nibabel.orientations`` used to resample to RAS 1 mm."""
-
-    def axcodes2ornt(self, axcodes: tuple[str, str, str]) -> HostArray: ...
-    def ornt_transform(self, start_ornt: HostArray, end_ornt: HostArray) -> HostArray: ...
-    def apply_orientation(self, data: HostArray, ornt: HostArray) -> HostArray: ...
-    def inv_ornt_aff(self, ornt: HostArray, shape: tuple[int, ...]) -> SpatialAffine: ...
-
-
-class Nibabel(Protocol):
-    Nifti1Header: HeaderFactory
-    Nifti1Image: ImageFactory
-    orientations: OrientationFunctions
-
-    def load(self, filename: PathArg) -> NiftiImage: ...
-    def save(self, image: NiftiImage, filename: PathArg) -> None: ...
-    def aff2axcodes(self, affine: HostArray) -> tuple[str | None, ...]: ...
-    def io_orientation(self, affine: HostArray) -> HostArray: ...
-
-
-# This adapter is used only for .nii/.nii.gz inputs and NIfTI output; the wider
-# NiBabel loader supports unrelated image classes outside the runner contract.
-nib = cast(Nibabel, nibabel)

lesiondiffusion_mlx-0.2.3/src/lesiondiffusion_mlx/_typing.py

@@ -1,122 +0,0 @@
-"""Native array and dependency contracts used by the source implementation."""
-
-from __future__ import annotations
-
-from typing import TYPE_CHECKING, Literal, Protocol
-
-from tokenizers import Encoding
-
-from lesiondiffusion_mlx._arrays import HostArray as HostArray
-
-if TYPE_CHECKING:
-    from collections.abc import Callable, Sequence
-
-    import mlx.core as mx
-    from medmlx_core.typing import FastRuntime as CoreFastRuntime
-    from medmlx_core.typing import MlxRuntime
-
-    type Array = mx.array
-    type Dtype = mx.Dtype
-    type ArrayValue = bool | int | float | complex | list[ArrayValue] | tuple[ArrayValue, ...]
-
-    class ArrayConstructor(Protocol):
-        def __call__(
-            self,
-            value: HostArray | Array | ArrayValue,
-            /,
-            dtype: Dtype | None = None,
-        ) -> Array: ...
-
-    class TokenizerContract(Protocol):
-        def enable_truncation(self, max_length: int) -> None: ...
-        def enable_padding(self, *, length: int, pad_id: int, pad_token: str) -> None: ...
-        def encode_batch(self, inputs: Sequence[str]) -> list[Encoding]: ...
-
-    class MetalKernel(Protocol):
-        def __call__(
-            self,
-            *,
-            inputs: Sequence[Array],
-            template: Sequence[tuple[str, int | bool | Dtype]],
-            grid: tuple[int, int, int],
-            threadgroup: tuple[int, int, int],
-            output_shapes: Sequence[Sequence[int]],
-            output_dtypes: Sequence[Dtype],
-        ) -> list[Array]: ...
-
-    class Fast(CoreFastRuntime, Protocol):
-        def scaled_dot_product_attention(
-            self,
-            q: Array,
-            k: Array,
-            v: Array,
-            *,
-            scale: float,
-            mask: str | Array | None = None,
-            sinks: Array | None = None,
-            force_fused: bool = False,
-        ) -> Array: ...
-
-        def metal_kernel(
-            self,
-            name: str,
-            input_names: Sequence[str],
-            output_names: Sequence[str],
-            source: str,
-            header: str = "",
-            ensure_row_contiguous: bool = True,
-            atomic_outputs: bool = False,
-            compile_options: object | None = None,
-        ) -> MetalKernel: ...
-
-    class Random(Protocol):
-        def seed(self, seed: int) -> None: ...
-        def normal(self, shape: Sequence[int]) -> Array: ...
-
-    class Mlx(MlxRuntime, Protocol):
-        """The MLX module surface consumed by this port, without importing it eagerly."""
-
-        @property
-        def fast(self) -> Fast: ...
-
-        random: Random
-
-        def addmm(self, c: Array, a: Array, b: Array, /) -> Array: ...
-        def all(self, a: Array, /) -> Array: ...
-        def argmin(self, a: Array, /, *, axis: int | None = None) -> Array: ...
-        def clear_cache(self) -> None: ...
-        def compile[**P, R](self, fun: Callable[P, R], /) -> Callable[P, R]: ...
-        def contiguous(self, a: Array, /) -> Array: ...
-        def cos(self, a: Array, /) -> Array: ...
-        def device_info(self) -> dict[str, str | int]: ...
-        def erf(self, a: Array, /) -> Array: ...
-        def exp(self, a: Array, /) -> Array: ...
-        def full(self, shape: Sequence[int], vals: object, *, dtype: Dtype) -> Array: ...
-        def isfinite(self, a: Array, /) -> Array: ...
-        def load(self, file: str, *, format: Literal["safetensors"]) -> dict[str, Array]: ...
-        def reciprocal(self, a: Array, /) -> Array: ...
-        def sin(self, a: Array, /) -> Array: ...
-        def softmax(
-            self, a: Array, /, *, axis: int | None = None, precise: bool = False
-        ) -> Array: ...
-        def sqrt(self, a: Array | float, /) -> Array: ...
-        def sum(self, a: Array, /, *, axis: int | None = None, keepdims: bool = False) -> Array: ...
-else:
-
-    class ArrayConstructor(Protocol):
-        pass
-
-    class TokenizerContract(Protocol):
-        pass
-
-    class MetalKernel(Protocol):
-        pass
-
-    class Fast(Protocol):
-        pass
-
-    class Random(Protocol):
-        pass
-
-    class Mlx(Protocol):
-        pass