diffract-core 0.2.1__py3-none-any.whl

diffract/core/data/nn/extractors/handlers/base.py

@@ -0,0 +1,110 @@
+"""Parameter handler interface and lightweight registry.
+
+This module defines the abstract ParameterHandler interface and a minimal
+registry for selecting and executing handlers. The design is intentionally
+simple: handlers are tried in registration order without priorities, and
+the registry returns both the processed weights and the handler instance
+that performed the processing.
+
+Key points:
+    - Single-pass selection in registration order
+    - Handlers declare produced ParameterType via parameter_type property
+    - Registry returns (ptype, weights, handler) to avoid double lookup
+    - Supported types are derived from registered handlers
+
+Example:
+    >>> registry = ParameterHandlerRegistry()
+    >>> registry.register_handler_class(MyDenseHandler)
+    >>> out = registry.process_parameter(module, "layer.weight")
+    >>> if out:
+    ...     ptype, weights, handler = out
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    import numpy as np
+    from numpy.typing import NDArray
+
+    from diffract.core.data.nn.params.schema import ParameterType
+
+
+class ParameterHandler[T](ABC):
+    """Base class for a concrete parameter-type handler (no priority)."""
+
+    @property
+    @abstractmethod
+    def parameter_type(self) -> ParameterType:
+        """ParameterType produced by this handler."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def can_handle(self, parameter: T, param_name: str) -> bool:
+        """Return True if handler can process the given parameter."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def process(self, parameter: T, param_name: str) -> NDArray[np.floating[Any]]:
+        """Return processed weights as a NumPy array."""
+        raise NotImplementedError
+
+    def get_additional_metadata(
+        self, _parameter: T, _param_name: str
+    ) -> dict[str, Any]:
+        """Return optional metadata; never raises and returns empty on failure."""
+        return {}
+
+
+class ParameterHandlerRegistry[T]:
+    """Registry of handlers with registration-order selection."""
+
+    def __init__(self) -> None:
+        self._handlers: list[ParameterHandler[T]] = []
+
+    def register_handler(self, handler: ParameterHandler[T]) -> None:
+        """Register a handler instance for later selection."""
+        self._handlers.append(handler)
+
+    def register_handler_class(self, handler_cls: type[ParameterHandler[T]]) -> None:
+        """Instantiate and register a handler class."""
+        self.register_handler(handler_cls())
+
+    def list_handlers(self) -> list[ParameterHandler[T]]:
+        """Return a shallow copy of registered handlers."""
+        return list(self._handlers)
+
+    def get_supported_types(self) -> list[ParameterType]:
+        """Return unique ParameterType values supported by registered handlers."""
+        seen: set[ParameterType] = set()
+        out: list[ParameterType] = []
+        for h in self._handlers:
+            if h.parameter_type not in seen:
+                seen.add(h.parameter_type)
+                out.append(h.parameter_type)
+        return out
+
+    def process_parameter(
+        self,
+        parameter: T,
+        param_name: str,
+    ) -> tuple[ParameterType, NDArray[np.floating[Any]], ParameterHandler[T]] | None:
+        """Select a matching handler and process in one pass.
+
+        Tries handlers in registration order, returning the first match's
+        produced type, processed weights, and the handler instance.
+
+        Args:
+            parameter: Framework-specific parameter object.
+            param_name: Human-readable parameter name.
+
+        Returns:
+            Tuple of (parameter type, weights array, handler) or None if no match.
+        """
+        for handler in self._handlers:
+            if handler.can_handle(parameter, param_name):
+                weights = handler.process(parameter, param_name)
+                return handler.parameter_type, weights, handler
+        return None

diffract/core/data/nn/extractors/handlers/flax_handlers.py

@@ -0,0 +1,71 @@
+"""Flax/JAX-specific parameter handlers for dense parameters.
+
+This module implements handlers for extracting dense layer weights from Flax
+parameter trees (typically `variables["params"]`).
+
+Currently supports only DENSE weights by selecting 2D arrays whose leaf name
+is `kernel` (Flax Linen Dense convention).
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any, Self, cast
+
+import diffract.core.utils.imports as import_utils
+from diffract.core.data.nn.params.schema import ParameterType
+
+from .base import ParameterHandler
+
+if TYPE_CHECKING:
+    import numpy as np
+    from numpy.typing import NDArray
+
+logger = logging.getLogger(__name__)
+
+
+if not (import_utils.is_available("flax") and import_utils.is_available("jax")):
+    logger.debug("Flax/JAX not available, disabling corresponding handlers")
+
+    class FlaxDenseHandler(ParameterHandler):
+        """Stub handler when Flax/JAX is not available."""
+
+        def __new__(  # noqa: D102
+            cls,
+            *_args: Any,
+            **_kwargs: Any,  # type: ignore[name-defined]
+        ) -> Self:
+            msg = "Flax/JAX packages not available"
+            raise ImportError(msg)
+
+else:
+    jax = import_utils.require("jax")
+    import numpy as np
+
+    class FlaxDenseHandler(ParameterHandler[Any]):
+        """Handler for Flax Dense kernel weights in params trees."""
+
+        _EXPECTED_NDIM = 2
+
+        @property
+        def parameter_type(self) -> ParameterType:  # noqa: D102
+            return ParameterType.DENSE
+
+        def can_handle(self, parameter: Any, param_name: str) -> bool:  # noqa: D102
+            # `parameter` is expected to be a leaf array-like.
+            if not param_name.endswith(".kernel"):
+                return False
+            arr = np.asarray(jax.device_get(parameter))
+            return arr.ndim == self._EXPECTED_NDIM and arr.size > 1
+
+        def process(  # noqa: D102
+            self, parameter: Any, _param_name: str
+        ) -> NDArray[np.floating[Any]]:
+            arr = np.asarray(jax.device_get(parameter))
+            if arr.ndim != self._EXPECTED_NDIM:
+                msg = (
+                    f"Expected {self._EXPECTED_NDIM}D kernel array, "
+                    f"got shape {arr.shape}"
+                )
+                raise ValueError(msg)
+            return cast("NDArray[np.floating[Any]]", arr)

diffract/core/data/nn/extractors/handlers/numpy_handlers.py

@@ -0,0 +1,63 @@
+"""NumPy-specific parameter handlers for dense parameters.
+
+This module implements handlers for extracting dense (fully-connected) layer
+weights from plain dictionaries of NumPy arrays. Handlers declare their
+produced ParameterType and provide optional metadata.
+
+NumPy is a required dependency, so no availability stubs are defined.
+
+Example:
+    >>> handler = NumpyDenseHandler()
+    >>> if handler.can_handle(array, name):
+    ...     weights = handler.process(array, name)
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, cast
+
+import numpy as np
+from numpy.typing import NDArray
+
+from diffract.core.data.nn.params.schema import ParameterType
+
+from .base import ParameterHandler
+
+logger = logging.getLogger(__name__)
+
+_DENSE_WEIGHT_NDIM = 2
+
+
+class NumpyDenseHandler(ParameterHandler[NDArray[Any]]):
+    """Handler for 2D floating-point arrays in NumPy dictionaries."""
+
+    @property
+    def parameter_type(self) -> ParameterType:
+        """Return produced type for this handler (DENSE)."""
+        return ParameterType.DENSE
+
+    def can_handle(self, parameter: NDArray[Any], _param_name: str) -> bool:
+        """Return True for 2D floating-point arrays with more than one element.
+
+        Masked arrays are rejected: their mask would be silently discarded,
+        computing metrics over raw fill values.
+        """
+        return (
+            parameter.ndim == _DENSE_WEIGHT_NDIM
+            and parameter.size > 1
+            and np.issubdtype(parameter.dtype, np.floating)
+            and not isinstance(parameter, np.ma.MaskedArray)
+        )
+
+    def process(
+        self, parameter: NDArray[Any], _param_name: str
+    ) -> NDArray[np.floating[Any]]:
+        """Return the array as NumPy weights."""
+        return cast("NDArray[np.floating[Any]]", parameter)
+
+    def get_additional_metadata(
+        self, parameter: NDArray[Any], _param_name: str
+    ) -> dict[str, Any]:
+        """Return optional NumPy-specific metadata such as dtype."""
+        return {"numpy_dtype": str(parameter.dtype)}

diffract/core/data/nn/extractors/handlers/onnx_handlers.py

@@ -0,0 +1,67 @@
+"""ONNX-specific parameter handlers for dense parameters.
+
+This module implements handlers for extracting dense weights from ONNX models.
+The extractor is responsible for selecting candidate weight arrays; this
+handler validates and standardizes them.
+
+Currently supports only DENSE weights by accepting 2D NumPy arrays.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any, Self, cast
+
+import diffract.core.utils.imports as import_utils
+from diffract.core.data.nn.params.schema import ParameterType
+
+from .base import ParameterHandler
+
+if TYPE_CHECKING:
+    import numpy as np
+    from numpy.typing import NDArray
+
+logger = logging.getLogger(__name__)
+
+
+if not import_utils.is_available("onnx"):
+    logger.debug("ONNX not available, disabling corresponding handlers")
+
+    class OnnxDenseHandler(ParameterHandler):
+        """Stub handler when ONNX is not available."""
+
+        def __new__(  # noqa: D102
+            cls,
+            *_args: Any,
+            **_kwargs: Any,  # type: ignore[name-defined]
+        ) -> Self:
+            msg = "onnx package not available"
+            raise ImportError(msg)
+
+else:
+    import numpy as np
+
+    class OnnxDenseHandler(ParameterHandler[Any]):
+        """Handler for ONNX dense weights represented as 2D NumPy arrays."""
+
+        _EXPECTED_NDIM = 2
+
+        @property
+        def parameter_type(self) -> ParameterType:  # noqa: D102
+            return ParameterType.DENSE
+
+        def can_handle(self, parameter: Any, _param_name: str) -> bool:  # noqa: D102
+            arr = np.asarray(parameter)
+            return arr.ndim == self._EXPECTED_NDIM and arr.size > 1
+
+        def process(  # noqa: D102
+            self, parameter: Any, _param_name: str
+        ) -> NDArray[np.floating[Any]]:
+            arr = np.asarray(parameter)
+            if arr.ndim != self._EXPECTED_NDIM:
+                msg = (
+                    f"Expected {self._EXPECTED_NDIM}D ONNX weight array, "
+                    f"got shape {arr.shape}"
+                )
+                raise ValueError(msg)
+            return cast("NDArray[np.floating[Any]]", arr)

diffract/core/data/nn/extractors/handlers/tensorflow_handlers.py

@@ -0,0 +1,82 @@
+"""TensorFlow-specific parameter handlers for dense parameters.
+
+This module implements handlers for extracting dense (fully-connected) layer
+weights from TensorFlow/Keras models. Currently supports only DENSE weights
+from `tf.keras.layers.Dense`.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any, Self, cast
+
+import diffract.core.utils.imports as import_utils
+from diffract.core.data.nn.params.schema import ParameterType
+
+from .base import ParameterHandler
+
+if TYPE_CHECKING:
+    import numpy as np
+    from numpy.typing import NDArray
+
+logger = logging.getLogger(__name__)
+
+
+if not import_utils.is_available("tensorflow"):
+    logger.debug("TensorFlow not available, disabling corresponding handlers")
+
+    class TensorFlowDenseHandler(ParameterHandler):
+        """Stub handler when TensorFlow is not available."""
+
+        def __new__(  # noqa: D102
+            cls,
+            *_args: Any,
+            **_kwargs: Any,  # type: ignore[name-defined]
+        ) -> Self:
+            msg = "TensorFlow package not available"
+            raise ImportError(msg)
+
+else:
+    tf = import_utils.require("tensorflow")
+
+    class TensorFlowDenseHandler(ParameterHandler[Any]):
+        """Handler for TensorFlow/Keras Dense layers (tf.keras.layers.Dense)."""
+
+        @property
+        def parameter_type(self) -> ParameterType:  # noqa: D102
+            return ParameterType.DENSE
+
+        def can_handle(self, parameter: Any, _param_name: str) -> bool:  # noqa: D102
+            return isinstance(parameter, tf.keras.layers.Dense)
+
+        def process(  # noqa: D102
+            self, parameter: Any, _param_name: str
+        ) -> NDArray[np.floating[Any]]:
+            if not isinstance(parameter, tf.keras.layers.Dense):
+                msg = f"Expected tf.keras.layers.Dense, got {type(parameter)}"
+                raise TypeError(msg)
+
+            # Keras Dense exposes kernel as a Variable. Ensure it's created/built.
+            if not hasattr(parameter, "kernel"):
+                msg = "Dense layer has no 'kernel' attribute (model not built?)"
+                raise ValueError(msg)
+
+            kernel = parameter.kernel
+            # In eager mode `.numpy()` is available; otherwise Keras models passed
+            # for analysis are expected to be materialized.
+            arr = kernel.numpy()
+            return cast("NDArray[np.floating[Any]]", arr)
+
+        def get_additional_metadata(  # noqa: D102
+            self, parameter: Any, _param_name: str
+        ) -> dict[str, Any]:
+            if not isinstance(parameter, tf.keras.layers.Dense):
+                return {}
+            meta: dict[str, Any] = {
+                "tf_dtype": getattr(parameter.kernel, "dtype", None),
+            }
+            # Useful for downstream analysis / sanity checks.
+            with_size = getattr(parameter, "units", None)
+            if with_size is not None:
+                meta["units"] = with_size
+            return meta

diffract/core/data/nn/extractors/handlers/torch_handlers.py

@@ -0,0 +1,124 @@
+"""PyTorch-specific parameter handlers for dense parameters.
+
+This module implements handlers for extracting dense (fully-connected) layer
+weights from PyTorch models and state dictionaries. Handlers declare their
+produced ParameterType and provide optional metadata.
+
+Behavior when PyTorch is unavailable:
+    - Stub handler classes are defined that raise ImportError on instantiation
+      to signal the missing dependency.
+
+Example:
+    >>> handler = TorchDenseHandler()
+    >>> if handler.can_handle(module, name):
+    ...     weights = handler.process(module, name)
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any, Self, cast
+
+import diffract.core.utils.imports as import_utils
+from diffract.core.data.nn.params.schema import ParameterType
+
+from .base import ParameterHandler
+
+if TYPE_CHECKING:
+    import numpy as np
+    from numpy.typing import NDArray
+
+logger = logging.getLogger(__name__)
+
+_DENSE_WEIGHT_NDIM = 2
+
+
+if not import_utils.is_available("torch"):
+    logger.debug("PyTorch not available, disabling corresponding handlers")
+
+    class TorchDenseHandler(ParameterHandler):
+        """Stub handler when PyTorch is not available."""
+
+        def __new__(
+            cls,
+            *_args: Any,
+            **_kwargs: Any,  # type: ignore[name-defined]
+        ) -> Self:
+            """Raise ImportError since PyTorch is unavailable."""
+            msg = "PyTorch package not available"
+            raise ImportError(msg)
+
+    class TorchStateDictDenseHandler(ParameterHandler):
+        """Stub handler when PyTorch is not available."""
+
+        def __new__(
+            cls,
+            *_args: Any,
+            **_kwargs: Any,  # type: ignore[name-defined]
+        ) -> Self:
+            """Raise ImportError since PyTorch is unavailable."""
+            msg = "PyTorch package not available"
+            raise ImportError(msg)
+
+else:
+    torch = import_utils.require("torch")
+    nn = torch.nn
+
+    class TorchDenseHandler(ParameterHandler[nn.Module]):
+        """Handler for PyTorch Linear/Dense layers (nn.Linear)."""
+
+        @property
+        def parameter_type(self) -> ParameterType:
+            """Return produced type for this handler (DENSE)."""
+            return ParameterType.DENSE
+
+        def can_handle(self, parameter: nn.Module, _param_name: str) -> bool:
+            """Return True if parameter is a PyTorch nn.Linear module."""
+            return isinstance(parameter, nn.Linear)
+
+        def process(
+            self, parameter: nn.Module, _param_name: str
+        ) -> NDArray[np.floating[Any]]:
+            """Convert nn.Linear weights to NumPy array using CPU tensors."""
+            if not isinstance(parameter, nn.Linear):
+                msg = f"Expected nn.Linear, got {type(parameter)}"
+                raise TypeError(msg)
+            return cast(
+                "NDArray[np.floating[Any]]", parameter.weight.detach().cpu().numpy()
+            )
+
+        def get_additional_metadata(
+            self, parameter: nn.Module, _param_name: str
+        ) -> dict[str, Any]:
+            """Return optional torch-specific metadata such as dtype."""
+            if not isinstance(parameter, nn.Linear):
+                return {}
+            return {"torch_dtype": str(parameter.weight.dtype)}
+
+    class TorchStateDictDenseHandler(ParameterHandler[torch.Tensor]):
+        """Handler for 2D weight tensors in PyTorch state dicts."""
+
+        @property
+        def parameter_type(self) -> ParameterType:
+            """Return produced type for this handler (DENSE)."""
+            return ParameterType.DENSE
+
+        def can_handle(self, parameter: torch.Tensor, param_name: str) -> bool:
+            """Return True for 2D weight tensors with positive number of elements."""
+            return (
+                param_name.endswith("weight")
+                and len(parameter.shape) == _DENSE_WEIGHT_NDIM
+                and parameter.numel() > 1
+            )
+
+        def process(
+            self, parameter: torch.Tensor, _param_name: str
+        ) -> NDArray[np.floating[Any]]:
+            """Convert torch.Tensor to NumPy array using CPU tensor."""
+            return cast("NDArray[np.floating[Any]]", parameter.detach().cpu().numpy())
+
+        def get_additional_metadata(
+            self, parameter: torch.Tensor, _param_name: str
+        ) -> dict[str, Any]:
+            """Return optional torch-specific metadata such as dtype."""
+            return {"torch_dtype": str(parameter.dtype)}

diffract/core/data/nn/extractors/interface.py

@@ -0,0 +1,56 @@
+"""Parameter extractor interface and protocols.
+
+This module defines the core protocol that all parameter extractor
+implementations must follow. It provides a framework-agnostic contract
+for parameter extraction operations from neural network models.
+
+The interface ensures consistent parameter extraction behavior across
+different deep learning frameworks while allowing framework-specific
+optimizations and parameter handling strategies.
+
+Example:
+    >>> extractor: IParameterExtractor = get_extractor(model)
+    >>> parameters = extractor.extract_parameters(storage, cache)
+    >>> print(f"Extracted {len(parameters)} parameters")
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Protocol, runtime_checkable
+
+if TYPE_CHECKING:
+    from diffract.core.data.nn.params.interface import IParameterRepository
+
+
+@runtime_checkable
+class IParameterExtractor(Protocol):
+    """Protocol defining the interface for neural network parameter extractors.
+
+    Provides a unified interface for extracting parameters from neural network
+    models across different deep learning frameworks. All extractor implementations
+    must implement this protocol to ensure consistent behavior.
+
+    The extraction process should traverse the model structure and extract all
+    relevant parameters along with their metadata, returning a standardized
+    parameter collection that can be used across the diffract system.
+    """
+
+    def extract_parameters(
+        self,
+        parameter_repository: IParameterRepository,
+    ) -> None:
+        """Extract all parameters from the neural network model.
+
+        Traverses the model structure and extracts all relevant weight parameters
+        along with their metadata. The extraction should be framework-specific
+        and add standardized ParameterDataProxy objects to the repository.
+
+        Args:
+            parameter_repository: Repository to store extracted parameters.
+
+        Raises:
+            NotImplementedError: If parameter type is not supported.
+            ValueError: If model structure is invalid.
+            RuntimeError: If extraction process fails.
+        """
+        ...

diffract/core/data/nn/extractors/numpy.py

@@ -0,0 +1,94 @@
+"""NumPy-specific parameter extractors.
+
+This module provides a parameter extractor for plain dictionaries mapping
+parameter names to NumPy arrays. It implements the BaseParameterExtractor
+interface without requiring any deep learning framework, since NumPy arrays
+already match the library's internal weight format.
+
+Supported Objects:
+    - dict[str, numpy.ndarray]: Mappings of parameter names to arrays
+
+Example:
+    >>> import numpy as np
+    >>> weights = {"encoder.weight": np.random.rand(10, 5)}
+    >>> extractor = NumpyDictExtractor(weights)
+    >>> parameters = extractor.extract_parameters(storage, cache)
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any
+
+from .base import BaseParameterExtractor
+from .handlers import NumpyDenseHandler
+
+if TYPE_CHECKING:
+    from collections.abc import Generator
+
+    from numpy.typing import NDArray
+
+    from .base import ExtractorOverrides
+    from .handlers.base import ParameterHandler
+
+logger = logging.getLogger(__name__)
+
+
+class NumpyDictExtractor(BaseParameterExtractor):
+    """NumPy dictionary parameter extractor with handler-based processing.
+
+    Extracts parameters from dictionaries mapping names to NumPy arrays by
+    iterating through entries and processing them with registered handlers.
+    This is the framework-free entry point: weight matrices loaded from
+    ``.npy``/``.npz`` files or via ``safetensors.numpy`` can be analyzed
+    without installing any deep learning framework.
+
+    Example:
+        >>> import numpy as np
+        >>> weights = {"encoder.weight": np.random.rand(10, 5)}
+        >>> extractor = NumpyDictExtractor(weights)
+        >>> params = extractor.extract_parameters(storage, cache)
+
+    Attributes:
+        model: Dictionary mapping parameter names to arrays.
+    """
+
+    def __init__(
+        self,
+        model: dict[str, NDArray[Any]],
+        overrides: ExtractorOverrides | None = None,
+        skip_not_implemented_types: bool = True,
+        custom_handlers: list[ParameterHandler] | None = None,
+    ) -> None:
+        """Initialize NumPy dictionary extractor.
+
+        Args:
+            model: Dictionary mapping parameter names to numpy.ndarray.
+            overrides: Optional extraction behavior overrides.
+            skip_not_implemented_types: Whether to skip unsupported arrays.
+            custom_handlers: Optional list of custom parameter handlers.
+        """
+        super().__init__(model, overrides, skip_not_implemented_types, custom_handlers)
+
+    def _register_default_handlers(self) -> None:
+        """Register default handlers for NumPy array dictionaries.
+
+        Registers the NumpyDenseHandler for processing dense weight matrices.
+        """
+        self.handler_registry.register_handler_class(NumpyDenseHandler)
+        logger.debug(
+            "Registered %d default handlers for NumPy arrays (DENSE only)",
+            len(self.handler_registry.list_handlers()),
+        )
+
+    def _iter_parameters(self) -> Generator[tuple[str, NDArray[Any]], None, None]:
+        """Iterate over arrays in the dictionary.
+
+        Yields all name-array pairs from the dictionary, allowing handlers
+        to process each array appropriately.
+
+        Yields:
+            Tuples of (parameter_name, array) for each entry in the dict.
+        """
+        arrays: dict[str, NDArray[Any]] = self.model
+        yield from arrays.items()