diffract-core 0.2.1__py3-none-any.whl

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

@@ -0,0 +1,108 @@
+"""ONNX parameter extractors.
+
+Provides a parameter extractor for ONNX `ModelProto` objects.
+Currently supports only dense (fully-connected) weights, extracted from `Gemm`
+nodes when their weight input is present in graph initializers.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any, Self
+
+import diffract.core.utils.imports as import_utils
+
+from .base import BaseParameterExtractor
+
+if TYPE_CHECKING:
+    from collections.abc import Generator
+
+    from .base import ExtractorOverrides
+    from .handlers.base import ParameterHandler
+
+logger = logging.getLogger(__name__)
+
+
+if not import_utils.is_available("onnx"):
+    logger.debug("ONNX not available, disabling ONNX extractors")
+
+    class OnnxModelExtractor(BaseParameterExtractor):
+        """Stub implementation when ONNX is not available."""
+
+        def __new__(cls, *_args: Any, **_kwargs: Any) -> Self:  # noqa: D102
+            msg = "onnx package not available"
+            raise ImportError(msg)
+
+else:
+    import numpy as np
+
+    onnx = import_utils.require("onnx")
+    numpy_helper = import_utils.require("onnx.numpy_helper")
+
+    from .handlers import OnnxDenseHandler
+
+    def _onnx_node_display_name(node: Any) -> str:
+        if getattr(node, "name", ""):
+            return node.name
+        outputs = list(getattr(node, "output", []) or [])
+        if outputs:
+            return outputs[0]
+        return "gemm"
+
+    class OnnxModelExtractor(BaseParameterExtractor):
+        """ONNX ModelProto parameter extractor (DENSE only)."""
+
+        def __init__(
+            self,
+            model: Any,
+            overrides: ExtractorOverrides | None = None,
+            skip_not_implemented_types: bool = True,
+            custom_handlers: list[ParameterHandler] | None = None,
+        ) -> None:
+            super().__init__(
+                model, overrides, skip_not_implemented_types, custom_handlers
+            )
+
+        def _register_default_handlers(self) -> None:
+            self.handler_registry.register_handler_class(OnnxDenseHandler)
+            logger.debug(
+                "Registered %d default handlers for ONNX models (DENSE only)",
+                len(self.handler_registry.list_handlers()),
+            )
+
+        def _iter_parameters(self) -> Generator[tuple[str, Any], None, None]:
+            """Yield dense weight arrays from Gemm nodes.
+
+            For `Gemm` node: inputs are typically [A, B, C]. We treat B as the
+            weight matrix, and ignore bias (C).
+            """
+            model = self.model
+            if not isinstance(model, onnx.ModelProto):
+                msg = f"OnnxModelExtractor expects onnx.ModelProto, got {type(model)}"
+                raise TypeError(msg)
+
+            graph = model.graph
+            initializer_by_name: dict[str, Any] = {
+                init.name: init for init in graph.initializer
+            }
+
+            min_gemm_inputs = 2
+            for node in graph.node:
+                if getattr(node, "op_type", None) != "Gemm":
+                    continue
+
+                inputs = list(getattr(node, "input", []) or [])
+                if len(inputs) < min_gemm_inputs:
+                    continue
+
+                weight_name = inputs[1]
+                init = initializer_by_name.get(weight_name)
+                if init is None:
+                    continue
+
+                weight_arr = numpy_helper.to_array(init)
+                # Ensure NumPy array for downstream handlers.
+                weight_arr = np.asarray(weight_arr)
+
+                node_name = _onnx_node_display_name(node)
+                yield f"{node_name}.weight", weight_arr

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

@@ -0,0 +1,99 @@
+"""TensorFlow/Keras parameter extractors.
+
+Provides a parameter extractor for TensorFlow Keras models.
+Currently supports only dense (fully-connected) layer kernels.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any, Self
+
+import diffract.core.utils.imports as import_utils
+
+from .base import BaseParameterExtractor
+
+if TYPE_CHECKING:
+    from collections.abc import Generator
+
+    from .base import ExtractorOverrides
+    from .handlers.base import ParameterHandler
+
+logger = logging.getLogger(__name__)
+
+
+if not import_utils.is_available("tensorflow"):
+    logger.debug("TensorFlow not available, disabling TensorFlow extractors")
+
+    class TensorFlowModelExtractor(BaseParameterExtractor):
+        """Stub implementation when TensorFlow is not available."""
+
+        def __new__(cls, *_args: Any, **_kwargs: Any) -> Self:  # noqa: D102
+            msg = "TensorFlow package not available"
+            raise ImportError(msg)
+
+else:
+    tf = import_utils.require("tensorflow")
+
+    from .handlers import TensorFlowDenseHandler
+
+    def _walk_layers(obj: Any) -> Generator[Any, None, None]:
+        """Yield layers contained in a Keras model/layer.
+
+        Uses public attributes (`layers`) when available and falls back to yielding
+        the object itself. This avoids relying on private APIs that may change
+        across TensorFlow/Keras versions.
+        """
+        if isinstance(obj, tf.keras.Model):
+            for layer in obj.layers:
+                yield from _walk_layers(layer)
+            return
+
+        if isinstance(obj, tf.keras.layers.Layer):
+            yield obj
+            # Some composite layers expose nested layers via `.layers`.
+            nested = getattr(obj, "layers", None)
+            if isinstance(nested, list):
+                for layer in nested:
+                    yield from _walk_layers(layer)
+            return
+
+    class TensorFlowModelExtractor(BaseParameterExtractor):
+        """TensorFlow Keras model parameter extractor (DENSE only)."""
+
+        def __init__(
+            self,
+            model: Any,
+            overrides: ExtractorOverrides | None = None,
+            skip_not_implemented_types: bool = True,
+            custom_handlers: list[ParameterHandler] | None = None,
+        ) -> None:
+            super().__init__(
+                model, overrides, skip_not_implemented_types, custom_handlers
+            )
+
+        def _register_default_handlers(self) -> None:
+            # Register only the DENSE handler.
+            self.handler_registry.register_handler_class(TensorFlowDenseHandler)
+            logger.debug(
+                "Registered %d default handlers for TensorFlow models (DENSE only)",
+                len(self.handler_registry.list_handlers()),
+            )
+
+        def _iter_parameters(self) -> Generator[tuple[str, Any], None, None]:
+            """Yield Dense layers as parameters.
+
+            We intentionally yield only Dense layers and name parameters as
+            `<layer.name>.kernel` to match weight-level naming convention.
+            """
+            model = self.model
+            if not isinstance(model, (tf.keras.Model, tf.keras.layers.Layer)):
+                msg = (
+                    "TensorFlowModelExtractor expects a tf.keras.Model or "
+                    f"tf.keras.layers.Layer, got {type(model)}"
+                )
+                raise TypeError(msg)
+
+            for layer in _walk_layers(model):
+                if isinstance(layer, tf.keras.layers.Dense):
+                    yield f"{layer.name}.kernel", layer

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

@@ -0,0 +1,204 @@
+"""PyTorch-specific parameter extractors.
+
+This module provides parameter extractors specifically designed for PyTorch
+models and state dictionaries. It implements the BaseParameterExtractor
+interface with PyTorch-specific parameter iteration and handler registration.
+
+Supported PyTorch Objects:
+    - torch.nn.Module: Complete PyTorch models with named modules
+    - dict[str, torch.Tensor]: PyTorch state dictionaries
+
+The extractors use handler-based processing to support different parameter
+types while maintaining extensibility for custom parameter handling. Currently
+focused on dense (fully-connected) layer parameters with room for expansion.
+
+Example:
+    >>> import torch
+    >>> model = torch.nn.Sequential(
+    ...     torch.nn.Linear(10, 5), torch.nn.ReLU(), torch.nn.Linear(5, 1)
+    ... )
+    >>> extractor = TorchModuleExtractor(model)
+    >>> parameters = extractor.extract_parameters(storage, cache)
+    >>> # For state dict:
+    >>> state_dict = model.state_dict()
+    >>> extractor = TorchStateDictExtractor(state_dict)
+    >>> parameters = extractor.extract_parameters(storage, cache)
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any, Self
+
+import diffract.core.utils.imports as import_utils
+
+from .base import BaseParameterExtractor
+
+if TYPE_CHECKING:
+    from collections.abc import Generator
+
+    import torch
+    from torch import nn
+
+    from .base import ExtractorOverrides
+    from .handlers.base import ParameterHandler
+
+logger = logging.getLogger(__name__)
+
+
+if not import_utils.is_available("torch"):
+    logger.debug("PyTorch not available, disabling PyTorch extractors")
+
+    class TorchModuleExtractor(BaseParameterExtractor):
+        """Stub implementation when PyTorch is not available.
+
+        Raises ImportError when instantiated to indicate PyTorch dependency
+        is missing.
+        """
+
+        def __new__(cls, *_args: Any, **_kwargs: Any) -> Self:
+            """Raise ImportError since PyTorch is unavailable."""
+            msg = "PyTorch package not available"
+            raise ImportError(msg)
+
+    class TorchStateDictExtractor(BaseParameterExtractor):
+        """Stub implementation when PyTorch is not available.
+
+        Raises ImportError when instantiated to indicate PyTorch dependency
+        is missing.
+        """
+
+        def __new__(cls, *_args: Any, **_kwargs: Any) -> Self:
+            """Raise ImportError since PyTorch is unavailable."""
+            msg = "PyTorch package not available"
+            raise ImportError(msg)
+
+else:
+    from .handlers import TorchDenseHandler, TorchStateDictDenseHandler
+
+    class TorchModuleExtractor(BaseParameterExtractor):
+        """PyTorch nn.Module parameter extractor with handler-based processing.
+
+        Extracts parameters from PyTorch nn.Module objects by iterating through
+        named modules and processing them with registered handlers. Currently
+        supports dense (fully-connected) layer parameters with extensible
+        architecture for additional parameter types.
+
+        Example:
+            >>> import torch.nn as nn
+            >>> model = nn.Sequential(nn.Linear(128, 64), nn.Linear(64, 10))
+            >>> extractor = TorchModuleExtractor(model)
+            >>> params = extractor.extract_parameters(storage, cache)
+
+        Attributes:
+            model: PyTorch nn.Module to extract parameters from.
+        """
+
+        def __init__(
+            self,
+            model: nn.Module,
+            overrides: ExtractorOverrides | None = None,
+            skip_not_implemented_types: bool = True,
+            custom_handlers: list[ParameterHandler] | None = None,
+        ) -> None:
+            """Initialize PyTorch module extractor.
+
+            Args:
+                model: PyTorch nn.Module to extract parameters from.
+                overrides: Optional extraction behavior overrides.
+                skip_not_implemented_types: Whether to skip unsupported modules.
+                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 PyTorch modules.
+
+            Registers the TorchDenseHandler for processing dense/linear layers.
+            Additional handlers can be registered for other layer types.
+            """
+            # Register only the DENSE handler that was originally supported
+            self.handler_registry.register_handler_class(TorchDenseHandler)
+            logger.debug(
+                "Registered %d default handlers for PyTorch modules (DENSE only)",
+                len(self.handler_registry.list_handlers()),
+            )
+
+        def _iter_parameters(self) -> Generator[tuple[str, nn.Module], None, None]:
+            """Iterate over named modules in the PyTorch model.
+
+            Yields all named modules in the model, allowing handlers to process
+            each module type appropriately. This includes leaf modules and
+            intermediate containers.
+
+            Yields:
+                Tuples of (module_name, module) for each module in the model.
+            """
+            yield from self.model.named_modules()
+
+    class TorchStateDictExtractor(BaseParameterExtractor):
+        """PyTorch state_dict parameter extractor with handler-based processing.
+
+        Extracts parameters from PyTorch state dictionaries by iterating through
+        tensor entries and processing them with registered handlers. State dicts
+        provide direct access to parameter tensors without module structure.
+
+        This extractor is useful for processing saved model weights, checkpoint
+        files, or when working with parameter tensors directly without the
+        full model architecture.
+
+        Example:
+            >>> import torch
+            >>> state_dict = torch.load("model.pt")
+            >>> extractor = TorchStateDictExtractor(state_dict)
+            >>> params = extractor.extract_parameters(storage, cache)
+
+        Attributes:
+            model: Dictionary mapping parameter names to tensors.
+        """
+
+        def __init__(
+            self,
+            model: dict[str, torch.Tensor],
+            overrides: ExtractorOverrides | None = None,
+            skip_not_implemented_types: bool = True,
+            custom_handlers: list[ParameterHandler] | None = None,
+        ) -> None:
+            """Initialize PyTorch state dict extractor.
+
+            Args:
+                model: Dictionary mapping parameter names to torch.Tensor.
+                overrides: Optional extraction behavior overrides.
+                skip_not_implemented_types: Whether to skip unsupported tensors.
+                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 PyTorch state dict tensors.
+
+            Registers the TorchStateDictDenseHandler for processing dense layer
+            parameter tensors from state dictionaries.
+            """
+            # Register only the DENSE handler that was originally supported
+            self.handler_registry.register_handler_class(TorchStateDictDenseHandler)
+            logger.debug(
+                "Registered %d default handlers for PyTorch state_dict (DENSE only)",
+                len(self.handler_registry.list_handlers()),
+            )
+
+        def _iter_parameters(self) -> Generator[tuple[str, torch.Tensor], None, None]:
+            """Iterate over parameter tensors in the state dictionary.
+
+            Yields all parameter name-tensor pairs from the state dictionary,
+            allowing handlers to process each tensor appropriately.
+
+            Yields:
+                Tuples of (parameter_name, tensor) for each entry in state_dict.
+            """
+            state_dict: dict[str, torch.Tensor] = self.model
+            yield from state_dict.items()

diffract/core/data/nn/params/__init__.py

@@ -0,0 +1,27 @@
+"""Parameter management module.
+
+This package contains the core parameter management components:
+- schema: Domain type definitions and enumerations
+- interface: Protocol definitions for repositories, views, and proxies
+- proxy: Parameter proxy implementation
+- view: Parameter view implementation
+"""
+
+from .interface import IParameterProxy, IParameterRepository, IParameterView
+from .metadata import ParameterMetadata
+from .proxy import ParameterDataProxy
+from .schema import FieldName, ParameterIndex, ParameterType, ParameterUID
+from .view import ParameterView
+
+__all__ = [
+    "FieldName",
+    "IParameterProxy",
+    "IParameterRepository",
+    "IParameterView",
+    "ParameterDataProxy",
+    "ParameterIndex",
+    "ParameterMetadata",
+    "ParameterType",
+    "ParameterUID",
+    "ParameterView",
+]