lesiondiffusion-mlx 0.2.1__py3-none-any.whl

lesiondiffusion_mlx/_typing.py

@@ -0,0 +1,110 @@
+"""Native array and dependency contracts used by the source implementation."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, 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]],
+            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, *args: object, **kwargs: object) -> 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 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 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) -> 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