regnets 0.1.0__py3-none-any.whl

_kernel_ops.py

@@ -0,0 +1,144 @@
+"""Packed tensor operations used by neural kernel descriptions."""
+
+from dataclasses import dataclass
+from math import isfinite
+
+import torch
+from torch import Tensor, nn
+from torch.nn import functional as F
+
+
+def positive_int(value: int, name: str) -> None:
+    """Отклонить размерности, смысл которых изменился бы при приведении типа."""
+    if type(value) is not int or value < 1:
+        raise ValueError(f"{name} must be a positive integer.")
+
+
+def boolean(value: bool, name: str) -> None:
+    """Требовать явно заданное логическое значение."""
+    if type(value) is not bool:
+        raise ValueError(f"{name} must be a boolean.")
+
+
+def positive_real(value: float, name: str) -> None:
+    """Проверить, что числовой параметр конечен и положителен."""
+    if type(value) not in (int, float) or not isfinite(value) or value <= 0:
+        raise ValueError(f"{name} must be positive and finite.")
+
+
+@dataclass(frozen=True)
+class KernelShape:
+    """Размерности и правила работы с параметрами для метода build ядра."""
+
+    in_features: int
+    out_features: int
+    positions: int
+    shared: bool
+    device: torch.device | str | None = None
+    dtype: torch.dtype | None = None
+
+    def __post_init__(self) -> None:
+        for name in ("in_features", "out_features", "positions"):
+            positive_int(getattr(self, name), name)
+        boolean(self.shared, "shared")
+
+
+class PackedLinear(nn.Module):
+    """Применить одно аффинное преобразование параллельно ко всем позициям модулей."""
+
+    def __init__(self, shape: KernelShape, bias: bool = True) -> None:
+        super().__init__()
+        self.in_features = shape.in_features
+        self.out_features = shape.out_features
+        self.positions = shape.positions
+        self.shared = shape.shared
+        prefix = () if shape.shared else (shape.positions,)
+        self.weight = nn.Parameter(
+            torch.empty(
+                (*prefix, shape.out_features, shape.in_features),
+                device=shape.device,
+                dtype=shape.dtype,
+            )
+        )
+        self.bias = (
+            nn.Parameter(
+                torch.empty(
+                    (*prefix, shape.out_features),
+                    device=shape.device,
+                    dtype=shape.dtype,
+                )
+            )
+            if bias
+            else None
+        )
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        """Инициализировать параметры по локальному числу входов независимо от числа позиций."""
+        bound = self.in_features**-0.5
+        nn.init.uniform_(self.weight, -bound, bound)
+        if self.bias is not None:
+            nn.init.uniform_(self.bias, -bound, bound)
+
+    def forward(self, x: Tensor) -> Tensor:
+        """Преобразовать [..., positions, in_features] без цикла по позициям."""
+        if self.shared:
+            return F.linear(x, self.weight, self.bias)
+        batch_shape = x.shape[:-2]
+        matrices = x.reshape(-1, self.positions, self.in_features).transpose(0, 1)
+        y = torch.bmm(matrices, self.weight.transpose(1, 2)).transpose(0, 1)
+        y = y.reshape(*batch_shape, self.positions, self.out_features)
+        return y if self.bias is None else y + self.bias
+
+    def extra_repr(self) -> str:
+        return (
+            f"in_features={self.in_features}, out_features={self.out_features}, "
+            f"positions={self.positions}, shared={self.shared}, bias={self.bias is not None}"
+        )
+
+
+class PackedNorm(nn.Module):
+    """Нормализовать локальные признаки независимо для позиций и элементов пакета."""
+
+    def __init__(
+        self, shape: KernelShape, *, rms: bool, eps: float, affine: bool, bias: bool
+    ) -> None:
+        super().__init__()
+        self.features = shape.in_features
+        self.rms = rms
+        self.eps = eps
+        dimensions = (
+            (shape.in_features,)
+            if shape.shared
+            else (shape.positions, shape.in_features)
+        )
+        self.weight = (
+            nn.Parameter(torch.ones(dimensions, device=shape.device, dtype=shape.dtype))
+            if affine
+            else None
+        )
+        self.bias = (
+            nn.Parameter(
+                torch.zeros(dimensions, device=shape.device, dtype=shape.dtype)
+            )
+            if affine and bias
+            else None
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        """Сохранить тип активации, если параметры имеют большую точность."""
+        if self.rms:
+            y = F.rms_norm(x, (self.features,), eps=self.eps)
+        else:
+            y = F.layer_norm(x, (self.features,), eps=self.eps)
+        if self.weight is not None:
+            y = y * self.weight.to(dtype=y.dtype)
+        if self.bias is not None:
+            y = y + self.bias.to(dtype=y.dtype)
+        return y
+
+    def extra_repr(self) -> str:
+        return (
+            f"features={self.features}, rms={self.rms}, eps={self.eps}, "
+            f"affine={self.weight is not None}, bias={self.bias is not None}"
+        )

kernels.py

@@ -0,0 +1,270 @@
+"""Immutable descriptions of vectorized neural kernels and their basic blocks."""
+
+from abc import ABC, abstractmethod
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass, replace
+from math import isfinite
+
+from torch import Tensor, nn
+
+from _kernel_ops import (
+    KernelShape as KernelShape,
+)
+from _kernel_ops import (
+    PackedLinear,
+    PackedNorm,
+    boolean,
+    positive_int,
+    positive_real,
+)
+
+type Activation = Callable[[], nn.Module] | None
+
+__all__ = [
+    "Kernel",
+    "KernelShape",
+    "Dense",
+    "MLP",
+    "Sequential",
+    "Identity",
+    "Dropout",
+    "LayerNorm",
+    "RMSNorm",
+    "GatedMLP",
+]
+
+
+def _activation(factory: Activation) -> nn.Module:
+    if factory is None:
+        return nn.Identity()
+    if not callable(factory):
+        raise ValueError("activation must be a factory such as nn.SiLU.")
+    module = factory()
+    if not isinstance(module, nn.Module):
+        raise ValueError("activation must create an nn.Module.")
+    if list(module.parameters()) or list(module.buffers()):
+        raise ValueError(
+            "Use a Kernel block for activations with parameters or buffers."
+        )
+    return module
+
+
+def _same_features(shape: KernelShape) -> None:
+    if shape.in_features != shape.out_features:
+        raise ValueError("This block preserves its number of features.")
+
+
+class Kernel(ABC):
+    """Описать операцию над [..., positions, features] с сохранением позиций.
+
+    Подклассы создают новые экземпляры nn.Module с параметрами согласно shape.shared.
+    Они могут смешивать локальные признаки, но не элементы пакета и не позиции.
+    """
+
+    def output_features(self, in_features: int, default: int) -> int:
+        """Определить размерность выхода, когда блок находится внутри Sequential."""
+        return default
+
+    @abstractmethod
+    def build(self, shape: KernelShape) -> nn.Module:
+        """Создать все регистрируемые параметры до создания оптимизатора."""
+
+
+@dataclass(frozen=True)
+class Dense(Kernel):
+    """Описать аффинное преобразование с необязательной выходной активацией.
+
+    Args:
+        out_features: Явная локальная размерность или размерность выхода слоя.
+        bias: Добавить обучаемое смещение.
+        activation: Фабрика активации без состояния, например nn.SiLU.
+    """
+
+    out_features: int | None = None
+    bias: bool = True
+    activation: Activation = None
+
+    def __post_init__(self) -> None:
+        if self.out_features is not None:
+            positive_int(self.out_features, "out_features")
+        boolean(self.bias, "bias")
+
+    def output_features(self, in_features: int, default: int) -> int:
+        return default if self.out_features is None else self.out_features
+
+    def build(self, shape: KernelShape) -> nn.Module:
+        if self.out_features is not None and self.out_features != shape.out_features:
+            raise ValueError("Dense out_features does not match the layer output.")
+        linear = PackedLinear(shape, self.bias)
+        return (
+            linear
+            if self.activation is None
+            else nn.Sequential(linear, _activation(self.activation))
+        )
+
+
+@dataclass(frozen=True)
+class MLP(Kernel):
+    """Описать локальный многослойный перцептрон с определяемыми размерностями входа и выхода."""
+
+    hidden: int | Sequence[int] = (64,)
+    activation: Activation = nn.SiLU
+    output_activation: Activation = None
+    bias: bool = True
+
+    def __post_init__(self) -> None:
+        if type(self.hidden) is int:
+            hidden = (self.hidden,)
+        elif isinstance(self.hidden, Sequence) and not isinstance(
+            self.hidden, (str, bytes)
+        ):
+            hidden = tuple(self.hidden)
+        else:
+            raise ValueError(
+                "hidden must be an integer or a sequence of hidden widths."
+            )
+        for width in hidden:
+            positive_int(width, "hidden width")
+        object.__setattr__(self, "hidden", hidden)
+        boolean(self.bias, "bias")
+
+    def build(self, shape: KernelShape) -> nn.Module:
+        hidden = (self.hidden,) if isinstance(self.hidden, int) else self.hidden
+        widths = (shape.in_features, *hidden, shape.out_features)
+        modules = []
+        for i, (source, target) in enumerate(zip(widths, widths[1:])):
+            activation = (
+                self.output_activation if i == len(widths) - 2 else self.activation
+            )
+            local = replace(shape, in_features=source, out_features=target)
+            modules.append(Dense(bias=self.bias, activation=activation).build(local))
+        return nn.Sequential(*modules)
+
+
+@dataclass(frozen=True, init=False)
+class Sequential(Kernel):
+    """Объединить блоки; блоки без изменения размерности сохраняют прежнюю локальную размерность."""
+
+    blocks: tuple[Kernel, ...]
+
+    def __init__(self, *blocks: Kernel) -> None:
+        if any(not isinstance(block, Kernel) for block in blocks):
+            raise ValueError("Sequential expects Kernel descriptions.")
+        object.__setattr__(self, "blocks", tuple(blocks))
+
+    def output_features(self, in_features: int, default: int) -> int:
+        for block in self.blocks:
+            in_features = block.output_features(in_features, default)
+        return in_features
+
+    def build(self, shape: KernelShape) -> nn.Module:
+        current = shape.in_features
+        modules = []
+        for block in self.blocks:
+            target = block.output_features(current, shape.out_features)
+            modules.append(
+                block.build(replace(shape, in_features=current, out_features=target))
+            )
+            current = target
+        if current != shape.out_features:
+            raise ValueError("Sequential output width does not match the layer output.")
+        return nn.Sequential(*modules)
+
+
+class _Preserving(Kernel):
+    def output_features(self, in_features: int, default: int) -> int:
+        return in_features
+
+
+@dataclass(frozen=True)
+class Identity(_Preserving):
+    """Сохранить локальные признаки без использования параметров."""
+
+    def build(self, shape: KernelShape) -> nn.Module:
+        _same_features(shape)
+        return nn.Identity()
+
+
+@dataclass(frozen=True)
+class Dropout(_Preserving):
+    """Исключать отдельные локальные признаки при обучении с независимыми масками."""
+
+    p: float = 0.1
+
+    def __post_init__(self) -> None:
+        if (
+            type(self.p) not in (int, float)
+            or not isfinite(self.p)
+            or not 0 <= self.p <= 1
+        ):
+            raise ValueError("Dropout p must be finite and between 0 and 1.")
+
+    def build(self, shape: KernelShape) -> nn.Module:
+        _same_features(shape)
+        return nn.Dropout(self.p)
+
+
+@dataclass(frozen=True)
+class LayerNorm(_Preserving):
+    """Нормализовать весь локальный вектор признаков, не смешивая позиции и элементы пакета."""
+
+    eps: float = 1e-5
+    affine: bool = True
+    bias: bool = True
+
+    def __post_init__(self) -> None:
+        positive_real(self.eps, "eps")
+        boolean(self.affine, "affine")
+        boolean(self.bias, "bias")
+
+    def build(self, shape: KernelShape) -> nn.Module:
+        _same_features(shape)
+        return PackedNorm(
+            shape, rms=False, eps=self.eps, affine=self.affine, bias=self.bias
+        )
+
+
+@dataclass(frozen=True)
+class RMSNorm(_Preserving):
+    """Нормализовать локальную среднеквадратичную величину с необязательным обучаемым масштабом."""
+
+    eps: float = 1e-6
+    affine: bool = True
+
+    def __post_init__(self) -> None:
+        positive_real(self.eps, "eps")
+        boolean(self.affine, "affine")
+
+    def build(self, shape: KernelShape) -> nn.Module:
+        _same_features(shape)
+        return PackedNorm(shape, rms=True, eps=self.eps, affine=self.affine, bias=False)
+
+
+class _Gated(nn.Module):
+    def __init__(
+        self, shape: KernelShape, hidden: int, activation: Activation, bias: bool
+    ):
+        super().__init__()
+        self.input = PackedLinear(replace(shape, out_features=2 * hidden), bias)
+        self.output = PackedLinear(replace(shape, in_features=hidden), bias)
+        self.activation = _activation(activation)
+
+    def forward(self, x: Tensor) -> Tensor:
+        gate, value = self.input(x).chunk(2, dim=-1)
+        return self.output(self.activation(gate) * value)
+
+
+@dataclass(frozen=True)
+class GatedMLP(Kernel):
+    """Вычислить проекции затвора и значения одной упакованной аффинной операцией."""
+
+    hidden: int = 64
+    activation: Activation = nn.SiLU
+    bias: bool = True
+
+    def __post_init__(self) -> None:
+        positive_int(self.hidden, "hidden")
+        boolean(self.bias, "bias")
+
+    def build(self, shape: KernelShape) -> nn.Module:
+        return _Gated(shape, self.hidden, self.activation, self.bias)

modular.py

@@ -0,0 +1,272 @@
+"""Self-similar PyTorch networks built from backend-independent topologies."""
+
+from collections.abc import Sequence
+from copy import deepcopy
+from dataclasses import dataclass, field
+from math import prod
+from typing import cast
+
+import torch
+from torch import Tensor, nn
+
+from _kernel_ops import KernelShape, PackedLinear, boolean, positive_int
+from kernels import Dense, Kernel
+from topology import Topology
+from topology.io import TopologySpec
+
+__all__ = ["ModularNet", "Layer", "Residual"]
+
+
+@dataclass(frozen=True)
+class Residual:
+    """Обойти один слой, при необходимости обучая локальную проекцию без смещения."""
+
+    project: bool = False
+
+    def __post_init__(self) -> None:
+        boolean(self.project, "project")
+
+
+@dataclass(frozen=True)
+class Layer:
+    """Описать ядро слоя, правила работы с параметрами и размеры векторов контактов.
+
+    Args:
+        kernel: Неизменяемое описание, по которому создаётся векторизованное ядро.
+        shared: Использовать общие параметры для позиций внутри слоя.
+        residual: Обойти слой, сохранив его порядок сгруппированных контактов.
+        in_size: Размер вектора входного контакта или None для скалярных контактов.
+        out_size: Размер вектора выходного контакта или размер входа, если не задан.
+    """
+
+    kernel: Kernel = field(default_factory=Dense)
+    shared: bool = True
+    residual: bool | Residual = False
+    in_size: int | None = field(default=None, kw_only=True)
+    out_size: int | None = field(default=None, kw_only=True)
+
+    def __post_init__(self) -> None:
+        if self.in_size is not None:
+            positive_int(self.in_size, "in_size")
+        if self.out_size is not None:
+            positive_int(self.out_size, "out_size")
+            if self.in_size is None:
+                raise ValueError(
+                    "out_size requires vector input contacts with an in_size."
+                )
+        if not isinstance(self.kernel, Kernel):
+            raise ValueError("kernel must be a Kernel description.")
+        boolean(self.shared, "shared")
+        if type(self.residual) is not bool and not isinstance(self.residual, Residual):
+            raise ValueError("residual must be a boolean or Residual description.")
+
+
+class _Layer(nn.Module):
+    def __init__(self, description: Layer, shape: KernelShape, *, identity_skip: bool):
+        super().__init__()
+        self.out_features = shape.out_features
+        project = (
+            isinstance(description.residual, Residual) and description.residual.project
+        )
+        if description.residual and not project and not identity_skip:
+            raise ValueError(
+                "Identity residual requires p == g and equal vector sizes; use Residual(project=True)."
+            )
+        self.kernel = description.kernel.build(shape)
+        if not isinstance(self.kernel, nn.Module):
+            raise ValueError("Kernel.build must return an nn.Module.")
+        self.skip: nn.Module | None = None
+        if project:
+            self.skip = PackedLinear(shape, bias=False)
+        elif description.residual:
+            self.skip = nn.Identity()
+
+    def forward(self, x: Tensor) -> Tensor:
+        y = self.kernel(x)
+        if not isinstance(y, Tensor) or y.shape != (*x.shape[:-1], self.out_features):
+            raise ValueError(
+                "Kernel returned an incompatible batch, position or feature shape."
+            )
+        return y if self.skip is None else self.skip(x) + y
+
+
+def _sizes(descriptions: tuple[Layer, ...]) -> tuple[tuple[int, ...], bool]:
+    """Определить и проверить общие поля контактов до создания ядер."""
+    previous = descriptions[0].in_size
+    scalar = previous is None
+    values = [1 if previous is None else previous]
+    for m, layer in enumerate(descriptions):
+        if m and previous != layer.in_size:
+            raise ValueError(
+                f"Layer {m - 1} out_size={previous} does not match "
+                f"Layer {m} in_size={layer.in_size}."
+            )
+        previous = layer.in_size if layer.out_size is None else layer.out_size
+        values.append(1 if previous is None else previous)
+    return tuple(values), scalar
+
+
+class ModularNet(nn.Module):
+    """Применить компактную топологию с упакованными ядрами и остаточными связями.
+
+    Остаточные связи необязательны.
+
+    Args:
+        topology: Соединения, не зависящие от PyTorch.
+        layers: Одно описание, повторяемое с новыми параметрами, или n описаний.
+            Размеры соседних контактов и скалярный или векторный режим должны совпадать.
+        device: Начальное устройство для параметров и индексов.
+        dtype: Начальный тип параметров.
+    """
+
+    def __init__(
+        self,
+        topology: Topology,
+        layers: Layer | Sequence[Layer],
+        *,
+        device: torch.device | str | None = None,
+        dtype: torch.dtype | None = None,
+    ) -> None:
+        """Применить компактную топологию с упакованными ядрами и остаточными связями.
+
+        Остаточные связи необязательны.
+
+        Args:
+            topology: Соединения, не зависящие от PyTorch.
+            layers: Одно описание, повторяемое с новыми параметрами, или n описаний.
+                Размеры соседних контактов и скалярный или векторный режим должны совпадать.
+            device: Начальное устройство для параметров и индексов.
+            dtype: Начальный тип параметров.
+        """
+        super().__init__()
+        if not isinstance(topology, Topology):
+            raise ValueError("topology must be a Topology object.")
+        self.topology = topology
+        n = topology.n_layers
+        if isinstance(layers, Layer):
+            descriptions = (layers,) * n
+        elif isinstance(layers, Sequence) and not isinstance(layers, (str, bytes)):
+            descriptions = tuple(layers)
+        else:
+            raise ValueError(
+                "layers must be a Layer or a sequence of Layer descriptions."
+            )
+        if len(descriptions) != n or any(
+            not isinstance(item, Layer) for item in descriptions
+        ):
+            raise ValueError(f"layers must contain {n} Layer descriptions.")
+        self.sizes, self.scalar = _sizes(descriptions)
+        self.n_inputs, self.n_outputs = prod(self.p), prod(self.g)
+        self.positions = tuple(topology.n_modules(m) for m in range(n))
+        self.layers = nn.ModuleList()
+        for m, description in enumerate(descriptions):
+            shape = KernelShape(
+                self.p[m] * self.sizes[m],
+                self.g[m] * self.sizes[m + 1],
+                self.positions[m],
+                description.shared,
+                device,
+                dtype,
+            )
+            try:
+                layer = _Layer(
+                    description,
+                    shape,
+                    identity_skip=self.p[m] == self.g[m]
+                    and self.sizes[m] == self.sizes[m + 1],
+                )
+            except (ValueError, TypeError) as error:
+                raise ValueError(f"Layer {m}: {error}") from error
+            self.layers.append(layer)
+        self._permutation_names = []
+        for m, indices in enumerate(topology.permutations()):
+            name = f"_permutation_{m}"
+            identity = all(i == value for i, value in enumerate(indices))
+            tensor = (
+                None
+                if identity
+                else torch.tensor(indices, dtype=torch.long, device=device)
+            )
+            self.register_buffer(name, tensor, persistent=False)
+            self._permutation_names.append(name)
+        self._configuration = self._describe(descriptions)
+
+    @property
+    def p(self) -> tuple[int, ...]:
+        """Вернуть числа входных контактов топологии по слоям."""
+        return self.topology.p
+
+    @property
+    def g(self) -> tuple[int, ...]:
+        """Вернуть числа выходных контактов топологии по слоям."""
+        return self.topology.g
+
+    def _describe(self, descriptions: tuple[Layer, ...]) -> dict:
+        """Описать топологию и действующие настройки каждого слоя."""
+        return {
+            "topology": TopologySpec("", self.topology).to_dict(),
+            "layers": [
+                {
+                    "kernel": [
+                        [
+                            f"{type(module).__module__}.{type(module).__qualname__}",
+                            module.extra_repr(),
+                        ]
+                        for module in cast(_Layer, self.layers[m]).kernel.modules()
+                    ],
+                    "shared": layer.shared,
+                    "residual": (
+                        "project"
+                        if isinstance(layer.residual, Residual)
+                        and layer.residual.project
+                        else "identity"
+                        if layer.residual
+                        else "none"
+                    ),
+                    "in_size": None if self.scalar else self.sizes[m],
+                    "out_size": None if self.scalar else self.sizes[m + 1],
+                }
+                for m, layer in enumerate(descriptions)
+            ],
+        }
+
+    def get_extra_state(self) -> dict:
+        """Сохранить метаданные соединений и построения вместе с зарегистрированными параметрами."""
+        return deepcopy(self._configuration)
+
+    def set_extra_state(self, state: dict) -> None:
+        """Запретить загрузку весов при другой конфигурации соединений или ядер."""
+        if state != self._configuration:
+            raise RuntimeError(
+                "Checkpoint topology or layer configuration does not match this network."
+            )
+
+    def _route(self, x: Tensor, m: int) -> Tensor:
+        indices = getattr(self, self._permutation_names[m])
+        return x if indices is None else x.index_select(-2, indices)
+
+    def forward(self, x: Tensor) -> Tensor:
+        """Сохранить начальные размерности пакета и применить все n+1 перестановок."""
+        if not isinstance(x, Tensor) or not x.is_floating_point():
+            raise ValueError("Input must be a floating-point tensor.")
+        if self.scalar:
+            if x.ndim < 1 or x.shape[-1] != self.n_inputs:
+                raise ValueError(f"Expected scalar input [..., {self.n_inputs}].")
+            x = x.unsqueeze(-1)
+        elif x.ndim < 2 or x.shape[-2:] != (self.n_inputs, self.sizes[0]):
+            raise ValueError(
+                f"Expected vector input [..., {self.n_inputs}, {self.sizes[0]}]."
+            )
+        batch_shape = x.shape[:-2]
+        for m, layer in enumerate(self.layers):
+            grouped = self._route(x, m).reshape(
+                *batch_shape, self.positions[m], self.p[m] * self.sizes[m]
+            )
+            x = layer(grouped).reshape(
+                *batch_shape, self.positions[m] * self.g[m], self.sizes[m + 1]
+            )
+        x = self._route(x, len(self.layers))
+        return x.squeeze(-1) if self.scalar else x
+
+    def extra_repr(self) -> str:
+        return f"n_inputs={self.n_inputs}, n_outputs={self.n_outputs}, sizes={self.sizes}, scalar={self.scalar}"