| """Semantic runtime weight identities shared by export and graph maintenance.""" |
|
|
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| from enum import Enum |
|
|
|
|
| class WeightLayout(Enum): |
| """Static-data ordering expected by the receiving XNNPACK operator.""" |
|
|
| IDENTITY = "identity" |
| CONVOLUTION = "oihw-to-ohwi" |
| DEPTHWISE_CONVOLUTION = "depthwise-oihw-to-mhwc" |
|
|
|
|
| def _valid_shape(shape: tuple[int, ...]) -> bool: |
| return bool(shape) and all( |
| isinstance(extent, int) and not isinstance(extent, bool) and extent > 0 for extent in shape |
| ) |
|
|
|
|
| @dataclass(frozen=True) |
| class CheckpointTensor: |
| """A complete checkpoint tensor, optionally reordered for XNNPACK.""" |
|
|
| source: str |
| source_shape: tuple[int, ...] |
| layout: WeightLayout |
|
|
| def __post_init__(self) -> None: |
| if not self.source or not _valid_shape(self.source_shape): |
| raise ValueError("invalid checkpoint tensor derivation") |
|
|
|
|
| @dataclass(frozen=True) |
| class SplitProjectionWeight: |
| """Channels extracted when the runtime evaluates one projection in two parts.""" |
|
|
| source: str |
| input_channel_slice: tuple[int, int] |
|
|
| def __post_init__(self) -> None: |
| begin, end = self.input_channel_slice |
| if ( |
| not self.source |
| or type(begin) is not int |
| or type(end) is not int |
| or not 0 <= begin < end |
| ): |
| raise ValueError("invalid split-projection weight derivation") |
|
|
|
|
| @dataclass(frozen=True) |
| class ZeroBias: |
| """A synthesized zero bias required by one split projection.""" |
|
|
| channels: int |
|
|
| def __post_init__(self) -> None: |
| if type(self.channels) is not int or self.channels <= 0: |
| raise ValueError("invalid zero-bias derivation") |
|
|
|
|
| type WeightDerivation = CheckpointTensor | SplitProjectionWeight | ZeroBias |
|
|
|
|
| def weight_derivation_name(derivation: WeightDerivation) -> str: |
| """Name a runtime weight by model meaning, never learned content.""" |
| if isinstance(derivation, CheckpointTensor): |
| return derivation.source |
| if isinstance(derivation, SplitProjectionWeight): |
| begin, end = derivation.input_channel_slice |
| return f"{derivation.source}[input:{begin}:{end}]" |
| return f"projection.zero_bias.{derivation.channels}" |
|
|
|
|
| @dataclass(frozen=True) |
| class WeightBinding: |
| """The stable semantic identity and byte extent of one runtime weight.""" |
|
|
| name: str |
| size: int |
| derivation: WeightDerivation |
|
|
| def __post_init__(self) -> None: |
| if ( |
| self.name != weight_derivation_name(self.derivation) |
| or not isinstance(self.size, int) |
| or isinstance(self.size, bool) |
| or self.size <= 0 |
| ): |
| raise ValueError("invalid runtime weight binding") |
|
|
|
|
| @dataclass(frozen=True) |
| class LoweredWeight: |
| """One named, layout-converted byte payload ready for a runtime slot.""" |
|
|
| binding: WeightBinding |
| payload: bytes |
|
|
| def __post_init__(self) -> None: |
| if self.binding.size != len(self.payload): |
| raise ValueError("lowered weight size disagrees with payload") |
|
|
