From 6bf4bf4f05c5cb7a9e57006bd0d924777c884111 Mon Sep 17 00:00:00 2001 From: Concode0 Date: Tue, 19 May 2026 22:41:04 +0900 Subject: [PATCH 1/4] refactor: move primitive layout dispatch into core runtime --- core/runtime/__init__.py | 8 ++ core/runtime/actions.py | 96 +++++++++++++++++++++++ core/runtime/layers.py | 120 +++++++++++++++++++++++++++++ layers/primitives/_utils.py | 88 --------------------- layers/primitives/linear.py | 12 +-- layers/primitives/multi_rotor.py | 74 ++++++------------ layers/primitives/normalization.py | 20 ++--- layers/primitives/projection.py | 31 +++----- layers/primitives/reflection.py | 40 ++++------ layers/primitives/rotor.py | 64 +++++---------- 10 files changed, 312 insertions(+), 241 deletions(-) create mode 100644 core/runtime/actions.py create mode 100644 core/runtime/layers.py diff --git a/core/runtime/__init__.py b/core/runtime/__init__.py index 0c108e6..c2d7dc1 100644 --- a/core/runtime/__init__.py +++ b/core/runtime/__init__.py @@ -8,8 +8,10 @@ """Runtime algebra hosts and dense reference operations.""" from .accessors import as_multivector, compact_values, grade_indices, hermitian_signs, materialize_dense, resolve_layout +from .actions import compact_multi_versor_action, compact_versor_action, versor_vector_matrix from .algebra import CliffordAlgebra from .context import AlgebraContext +from .layers import LayerStorage, resolve_layer_layout, resolve_layer_storage from .multivector import Multivector from .projected import AlgebraRuntimeMixin @@ -20,8 +22,14 @@ "Multivector", "as_multivector", "compact_values", + "compact_multi_versor_action", + "compact_versor_action", "grade_indices", "hermitian_signs", "materialize_dense", + "resolve_layer_layout", + "resolve_layer_storage", "resolve_layout", + "LayerStorage", + "versor_vector_matrix", ] diff --git a/core/runtime/actions.py b/core/runtime/actions.py new file mode 100644 index 0000000..ba7b8c0 --- /dev/null +++ b/core/runtime/actions.py @@ -0,0 +1,96 @@ +"""Runtime actions used by compact-capable primitive layers.""" + +from __future__ import annotations + +import torch + +from core.foundation.layout import GradeLayout +from core.planning.action import bivector_vector_generator, metric_self_signs, reflection_vector_matrix +from core.runtime.accessors import materialize_dense + + +def compact_versor_action( + algebra, + values: torch.Tensor, + weights: torch.Tensor, + *, + grade: int, + input_layout: GradeLayout, + output_layout: GradeLayout, + parameter_layout: GradeLayout, + compact_output: bool, +) -> torch.Tensor: + """Apply one compact versor action to layer values.""" + matrix = versor_vector_matrix( + algebra, + weights.to(device=values.device, dtype=values.dtype), + grade=grade, + parameter_layout=parameter_layout, + ) + return algebra.planned_linear_action( + values, + matrix, + input_layout=input_layout, + output_layout=output_layout, + input_compact=True, + compact_output=compact_output, + ) + + +def compact_multi_versor_action( + algebra, + values: torch.Tensor, + weights: torch.Tensor, + mix: torch.Tensor, + *, + grade: int, + input_layout: GradeLayout, + output_layout: GradeLayout, + parameter_layout: GradeLayout, + compact_output: bool, +) -> torch.Tensor: + """Apply a weighted compact superposition of versor actions.""" + matrices = versor_vector_matrix( + algebra, + weights.to(device=values.device, dtype=values.dtype), + grade=grade, + parameter_layout=parameter_layout, + ) + outputs = [] + for index in range(matrices.shape[0]): + matrix = matrices[index].unsqueeze(0).expand(values.shape[-2], -1, -1) + outputs.append( + algebra.planned_linear_action( + values, + matrix, + input_layout=input_layout, + output_layout=output_layout, + input_compact=True, + compact_output=True, + ) + ) + + stacked = torch.stack(outputs, dim=-2) + result = torch.einsum("ck,...ckd->...cd", mix.to(device=values.device, dtype=values.dtype), stacked) + if compact_output: + return result + return materialize_dense(algebra, result, layout=output_layout) + + +def versor_vector_matrix(algebra, weights: torch.Tensor, *, grade: int, parameter_layout: GradeLayout) -> torch.Tensor: + """Return the vector-space matrix represented by compact versor weights.""" + grade = int(grade) + if grade == 2: + return torch.matrix_exp(bivector_vector_generator(weights, bivector_layout=parameter_layout)) + if grade == 1: + signs = parameter_layout_signs(parameter_layout, device=weights.device, dtype=weights.dtype) + norm_sq = (weights * weights * signs).sum(dim=-1, keepdim=True) + scale = norm_sq.abs().clamp_min(1e-12).sqrt() + normals = weights / scale + return reflection_vector_matrix(normals, vector_layout=parameter_layout, eps=algebra.eps_sq) + raise ValueError("compact versor execution currently supports grade=1 and grade=2") + + +def parameter_layout_signs(layout: GradeLayout, *, device=None, dtype=None) -> torch.Tensor: + """Return basis self-product signs for compact parameter weights.""" + return metric_self_signs(layout, device=device, dtype=dtype) diff --git a/core/runtime/layers.py b/core/runtime/layers.py new file mode 100644 index 0000000..2bfa4ad --- /dev/null +++ b/core/runtime/layers.py @@ -0,0 +1,120 @@ +"""Runtime helpers for layer-facing multivector storage.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from core.foundation.layout import GradeLayout +from core.planning.action import metric_self_signs + + +@dataclass(frozen=True) +class LayerStorage: + """Resolved storage contract for a layer input or output.""" + + algebra: object + layout: GradeLayout | None = None + + @property + def lane_dim(self) -> int: + """Return the coefficient lane count accepted by this storage.""" + return self.algebra.dim if self.layout is None else self.layout.dim + + @property + def is_compact(self) -> bool: + """Return whether this storage is compact relative to the full algebra.""" + return self.layout is not None and self.layout.dim != self.algebra.dim + + @property + def grades(self) -> tuple[int, ...] | None: + """Return active grades when compact metadata is known.""" + return None if self.layout is None else self.layout.grades + + def validate_input( + self, + values: torch.Tensor, + *, + channels: int, + name: str, + allow_dense: bool | None = None, + ) -> bool: + """Validate layer input and return whether it is compact.""" + if values.ndim < 3: + raise ValueError(f"{name}: expected ndim >= 3, got shape {tuple(values.shape)}") + if values.shape[-2] != channels: + raise ValueError( + f"{name}: expected {channels} channels, got {values.shape[-2]} (shape {tuple(values.shape)})" + ) + + if self.layout is not None and values.shape[-1] == self.layout.dim: + return self.is_compact + + if allow_dense is None: + allow_dense = self.layout is None or self.layout.dim == self.algebra.dim + if allow_dense and values.shape[-1] == self.algebra.dim: + return False + + expected = [str(self.algebra.dim)] if allow_dense else [] + if self.layout is not None: + expected.insert(0, f"{self.layout.dim} for grades {self.layout.grades}") + raise ValueError(f"{name}: last dim must be {' or '.join(expected)}, got {values.shape[-1]}") + + def scalar_mask(self, *, device=None, dtype=None) -> torch.Tensor: + """Return a scalar-lane mask for this storage.""" + dtype = torch.float32 if dtype is None else dtype + if self.layout is None: + mask = torch.zeros(self.algebra.dim, device=device, dtype=dtype) + mask[0] = 1.0 + return mask + return torch.tensor( + [1.0 if index == 0 else 0.0 for index in self.layout.basis_indices], + device=device, + dtype=dtype, + ) + + def grade_positions(self, grade: int, *, device=None) -> torch.Tensor: + """Return compact lane positions for one grade.""" + if self.layout is None: + return self.algebra.grade_indices((grade,), device=device) + positions = [ + position for position, index in enumerate(self.layout.basis_indices) if index.bit_count() == int(grade) + ] + return torch.tensor(positions, dtype=torch.long, device=device) + + def compact_grade_norms(self, values: torch.Tensor) -> torch.Tensor: + """Return per-grade coefficient norms for compact values.""" + if self.layout is None: + return self.algebra.get_grade_norms(values) + flat = values.pow(2).reshape(-1, self.layout.dim) + grade_ids = self.layout.grade_indices_tensor(device=values.device).unsqueeze(0).expand_as(flat) + result = values.new_zeros(flat.shape[0], self.algebra.num_grades) + result.scatter_add_(1, grade_ids, flat) + return result.reshape(*values.shape[:-1], self.algebra.num_grades).clamp(min=self.algebra.eps).sqrt() + + def metric_signs(self, *, device=None, dtype=None) -> torch.Tensor: + """Return basis self-product signs for this storage.""" + if self.layout is None: + return metric_self_signs(self.algebra.default_layout(), device=device, dtype=dtype) + return metric_self_signs(self.layout, device=device, dtype=dtype) + + +def resolve_layer_storage(algebra, *, layout: GradeLayout = None, grades=None) -> LayerStorage: + """Resolve optional layer grade/layout metadata into a storage contract.""" + return LayerStorage(algebra, resolve_layer_layout(algebra, layout=layout, grades=grades)) + + +def resolve_layer_layout(algebra, *, layout: GradeLayout = None, grades=None) -> GradeLayout | None: + """Resolve an optional layer storage layout.""" + if layout is not None: + spec = algebra.planner.spec + if layout.spec != spec: + raise ValueError(f"layout signature {layout.spec} does not match algebra signature {spec}") + return layout + if grades is not None: + return algebra.layout(grades) + default_grades = getattr(algebra, "_default_grades", None) + if default_grades is not None: + return algebra.layout(default_grades) + return None diff --git a/layers/primitives/_utils.py b/layers/primitives/_utils.py index c90abe6..cfe70b7 100644 --- a/layers/primitives/_utils.py +++ b/layers/primitives/_utils.py @@ -6,9 +6,6 @@ import torch -from core.foundation.layout import GradeLayout -from core.planning.action import metric_self_signs - def require_positive_int(value: int, name: str) -> int: """Validate a positive integer layer dimension.""" @@ -38,91 +35,6 @@ def grade_indices(algebra, grade: int, *, name: str = "grade") -> torch.Tensor: return indices -def resolve_layer_layout(algebra, *, layout: GradeLayout = None, grades=None) -> GradeLayout | None: - """Resolve an optional primitive storage layout.""" - if layout is not None: - spec = algebra.planner.spec - if layout.spec != spec: - raise ValueError(f"layout signature {layout.spec} does not match algebra signature {spec}") - return layout - if grades is not None: - return algebra.layout(grades) - default_grades = getattr(algebra, "_default_grades", None) - if default_grades is not None: - return algebra.layout(default_grades) - return None - - -def lane_dim(algebra, layout: GradeLayout | None) -> int: - """Return the active lane count for dense or compact storage.""" - return algebra.dim if layout is None else layout.dim - - -def is_compact_layout(algebra, layout: GradeLayout | None) -> bool: - """Return whether a resolved layout is strictly compact.""" - return layout is not None and layout.dim != algebra.dim - - -def check_multivector_storage( - x: torch.Tensor, - algebra, - *, - channels: int, - name: str, - layout: GradeLayout | None = None, - allow_dense: bool = True, -) -> bool: - """Validate primitive input and return whether it is compact.""" - if x.ndim < 3: - raise ValueError(f"{name}: expected ndim >= 3, got shape {tuple(x.shape)}") - if x.shape[-2] != channels: - raise ValueError(f"{name}: expected {channels} channels, got {x.shape[-2]} (shape {tuple(x.shape)})") - - if layout is not None and x.shape[-1] == layout.dim: - return layout.dim != algebra.dim - if allow_dense and x.shape[-1] == algebra.dim: - return False - - expected = [str(algebra.dim)] if allow_dense else [] - if layout is not None: - expected.insert(0, f"{layout.dim} for grades {layout.grades}") - raise ValueError(f"{name}: last dim must be {' or '.join(expected)}, got {x.shape[-1]}") - - -def scalar_mask(algebra, layout: GradeLayout | None, *, device=None, dtype=None) -> torch.Tensor: - """Return a scalar-lane mask for dense or compact storage.""" - if layout is None: - mask = torch.zeros(algebra.dim, device=device, dtype=torch.float32 if dtype is None else dtype) - mask[0] = 1.0 - return mask - values = torch.tensor( - [1.0 if index == 0 else 0.0 for index in layout.basis_indices], - device=device, - dtype=torch.float32 if dtype is None else dtype, - ) - return values - - -def grade_positions(layout: GradeLayout, grade: int) -> torch.Tensor: - """Return compact positions for one grade within a layout.""" - positions = [position for position, index in enumerate(layout.basis_indices) if index.bit_count() == int(grade)] - return torch.tensor(positions, dtype=torch.long) - - -def layout_metric_signs(layout: GradeLayout, *, device=None, dtype=None) -> torch.Tensor: - """Return basis self-product signs for a layer layout.""" - return metric_self_signs(layout, device=device, dtype=dtype) - - -def compact_grade_norms(algebra, values: torch.Tensor, layout: GradeLayout) -> torch.Tensor: - """Return per-grade coefficient norms for compact values.""" - flat = values.pow(2).reshape(-1, layout.dim) - grade_ids = layout.grade_indices_tensor(device=values.device).unsqueeze(0).expand_as(flat) - result = values.new_zeros(flat.shape[0], algebra.num_grades) - result.scatter_add_(1, grade_ids, flat) - return result.reshape(*values.shape[:-1], algebra.num_grades).clamp(min=algebra.eps).sqrt() - - def dense_from_indices(coefficients: torch.Tensor, indices: torch.Tensor, dense_dim: int) -> torch.Tensor: """Scatter compact coefficients into dense multivector storage.""" dense = coefficients.new_zeros(*coefficients.shape[:-1], dense_dim) diff --git a/layers/primitives/linear.py b/layers/primitives/linear.py index 3552ca3..41d0d86 100644 --- a/layers/primitives/linear.py +++ b/layers/primitives/linear.py @@ -18,8 +18,9 @@ from core.foundation.layout import GradeLayout from core.foundation.module import CliffordModule from core.runtime.algebra import CliffordAlgebra +from core.runtime.layers import resolve_layer_storage -from ._utils import check_multivector_storage, lane_dim, require_choice, require_positive_int, resolve_layer_layout +from ._utils import require_choice, require_positive_int class CliffordLinear(CliffordModule): @@ -73,8 +74,9 @@ def __init__( self.in_channels = require_positive_int(in_channels, "in_channels") self.out_channels = require_positive_int(out_channels, "out_channels") self.backend = require_choice(backend, "backend", ("traditional", "rotor")) - self.layout = resolve_layer_layout(algebra, layout=layout, grades=grades) - self.lane_dim = lane_dim(algebra, self.layout) + self.storage = resolve_layer_storage(algebra, layout=layout, grades=grades) + self.layout = self.storage.layout + self.lane_dim = self.storage.lane_dim if self.backend == "traditional": self.weight = nn.Parameter(torch.Tensor(self.out_channels, self.in_channels)) @@ -116,12 +118,10 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: Returns: torch.Tensor: Output [Batch, Out, Dim]. """ - check_multivector_storage( + self.storage.validate_input( x, - self.algebra, channels=self.in_channels, name="CliffordLinear input", - layout=self.layout, allow_dense=self.layout is None or self.layout.dim == self.algebra.dim, ) diff --git a/layers/primitives/multi_rotor.py b/layers/primitives/multi_rotor.py index 72ddcd2..38383ed 100644 --- a/layers/primitives/multi_rotor.py +++ b/layers/primitives/multi_rotor.py @@ -16,19 +16,15 @@ from core.foundation.layout import GradeLayout from core.foundation.manifold import MANIFOLD_SPIN, tag_manifold from core.foundation.module import CliffordModule -from core.planning.action import bivector_vector_generator, reflection_vector_matrix -from core.runtime.accessors import materialize_dense +from core.runtime.actions import compact_multi_versor_action from core.runtime.algebra import CliffordAlgebra +from core.runtime.layers import resolve_layer_storage from ._utils import ( cache_matches, - check_multivector_storage, - compact_grade_norms, dense_from_indices, grade_indices, - layout_metric_signs, require_positive_int, - resolve_layer_layout, ) @@ -76,12 +72,14 @@ def __init__( self.channels = require_positive_int(channels, "channels") self.num_rotors = require_positive_int(num_rotors, "num_rotors") self.grade = int(grade) - self.input_layout = resolve_layer_layout(algebra, layout=input_layout, grades=input_grades) - self.output_layout = ( - resolve_layer_layout(algebra, layout=output_layout, grades=output_grades) + self.input_storage = resolve_layer_storage(algebra, layout=input_layout, grades=input_grades) + self.output_storage = ( + resolve_layer_storage(algebra, layout=output_layout, grades=output_grades) if output_layout is not None or output_grades is not None - else self.input_layout + else self.input_storage ) + self.input_layout = self.input_storage.layout + self.output_layout = self.output_storage.layout self.compact_output = bool(compact_output) self.register_buffer("grade_indices", grade_indices(algebra, self.grade)) @@ -154,19 +152,18 @@ def forward(self, x: torch.Tensor, return_invariants: bool = False) -> torch.Ten Returns: torch.Tensor: Transformed output [Batch, Channels, Dim]. """ - is_compact = check_multivector_storage( + is_compact = self.input_storage.validate_input( x, - self.algebra, channels=self.channels, name="MultiRotorLayer input", - layout=self.input_layout, allow_dense=self.input_layout is None or self.input_layout.dim == self.algebra.dim, ) if is_compact: out = self._forward_compact(x) if return_invariants: - output_layout = self.output_layout if self.output_layout is not None else self.input_layout - return compact_grade_norms(self.algebra, out, output_layout) + if not self.compact_output: + return self.algebra.get_grade_norms(out) + return self.output_storage.compact_grade_norms(out) return out if not hasattr(self.algebra, "multi_rotor_sandwich"): raise ValueError( @@ -206,40 +203,19 @@ def _forward_compact(self, x: torch.Tensor) -> torch.Tensor: """Apply compact weighted superposition of induced versor actions.""" if self.input_layout is None: raise ValueError("MultiRotorLayer compact input requires input_layout or input_grades") - output_layout = self.output_layout if self.output_layout is not None else self.input_layout - - rotor_values = self.rotor_grade_weights.to(device=x.device, dtype=x.dtype) - if self.grade == 2: - generator = bivector_vector_generator(rotor_values, bivector_layout=self.parameter_layout) - matrices = torch.matrix_exp(generator) - elif self.grade == 1: - signs = layout_metric_signs(self.parameter_layout, device=x.device, dtype=x.dtype) - norm_sq = (rotor_values * rotor_values * signs).sum(dim=-1, keepdim=True) - scale = norm_sq.abs().clamp_min(1e-12).sqrt() - normals = rotor_values / scale - matrices = reflection_vector_matrix(normals, vector_layout=self.parameter_layout, eps=self.algebra.eps_sq) - else: - raise ValueError("MultiRotorLayer compact execution currently supports grade=1 and grade=2 versors") - - outputs = [] - for rotor_index in range(self.num_rotors): - matrix = matrices[rotor_index].unsqueeze(0).expand(self.channels, -1, -1) - outputs.append( - self.algebra.planned_linear_action( - x, - matrix, - input_layout=self.input_layout, - output_layout=output_layout, - input_compact=True, - compact_output=True, - ) - ) - stacked = torch.stack(outputs, dim=-2) - weights = self.weights.to(device=x.device, dtype=x.dtype) - out = torch.einsum("ck,...ckd->...cd", weights, stacked) - if self.compact_output: - return out - return materialize_dense(self.algebra, out, layout=output_layout) + if self.output_layout is None: + raise ValueError("MultiRotorLayer compact output requires output_layout or output_grades") + return compact_multi_versor_action( + self.algebra, + x, + self.rotor_grade_weights, + self.weights, + grade=self.grade, + input_layout=self.input_layout, + output_layout=self.output_layout, + parameter_layout=self.parameter_layout, + compact_output=self.compact_output, + ) def train(self, mode: bool = True): """Invalidate versor cache when switching to train mode.""" diff --git a/layers/primitives/normalization.py b/layers/primitives/normalization.py index 438922b..c6a0820 100644 --- a/layers/primitives/normalization.py +++ b/layers/primitives/normalization.py @@ -11,14 +11,9 @@ from core.foundation.layout import GradeLayout from core.foundation.module import CliffordModule from core.runtime.algebra import CliffordAlgebra +from core.runtime.layers import resolve_layer_storage -from ._utils import ( - check_multivector_storage, - lane_dim, - require_positive_int, - resolve_layer_layout, - scalar_mask, -) +from ._utils import require_positive_int class CliffordLayerNorm(CliffordModule): @@ -60,12 +55,13 @@ def __init__( raise ValueError(f"eps must be positive, got {eps}") self.eps = eps self.recover = recover - self.layout = resolve_layer_layout(algebra, layout=layout, grades=grades) - self.lane_dim = lane_dim(algebra, self.layout) + self.storage = resolve_layer_storage(algebra, layout=layout, grades=grades) + self.layout = self.storage.layout + self.lane_dim = self.storage.lane_dim self.weight = nn.Parameter(torch.ones(self.channels)) self.bias = nn.Parameter(torch.zeros(self.channels)) - self.register_buffer("scalar_mask", scalar_mask(algebra, self.layout)) + self.register_buffer("scalar_mask", self.storage.scalar_mask()) if recover: self.norm_scale = nn.Parameter(torch.zeros(self.channels)) else: @@ -80,12 +76,10 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: Returns: torch.Tensor: Normalized input. """ - check_multivector_storage( + self.storage.validate_input( x, - self.algebra, channels=self.channels, name="CliffordLayerNorm input", - layout=self.layout, allow_dense=self.layout is None or self.layout.dim == self.algebra.dim, ) channel_shape = (1,) * (x.ndim - 2) + (self.channels, 1) diff --git a/layers/primitives/projection.py b/layers/primitives/projection.py index ab9a184..c7de8e5 100644 --- a/layers/primitives/projection.py +++ b/layers/primitives/projection.py @@ -11,15 +11,10 @@ from core.foundation.layout import GradeLayout from core.foundation.module import CliffordModule from core.runtime.algebra import CliffordAlgebra +from core.runtime.layers import resolve_layer_storage from utils.compat import safe_linalg_solve -from ._utils import ( - check_multivector_storage, - grade_positions, - lane_dim, - require_positive_int, - resolve_layer_layout, -) +from ._utils import require_positive_int class BladeSelector(CliffordModule): @@ -40,8 +35,9 @@ def __init__(self, algebra: CliffordAlgebra, channels: int, *, grades=None, layo """ super().__init__(algebra) self.channels = require_positive_int(channels, "channels") - self.layout = resolve_layer_layout(algebra, layout=layout, grades=grades) - self.lane_dim = lane_dim(algebra, self.layout) + self.storage = resolve_layer_storage(algebra, layout=layout, grades=grades) + self.layout = self.storage.layout + self.lane_dim = self.storage.lane_dim self.weights = nn.Parameter(torch.Tensor(self.channels, self.lane_dim)) @@ -62,12 +58,10 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: Returns: torch.Tensor: Filtered input. """ - check_multivector_storage( + self.storage.validate_input( x, - self.algebra, channels=self.channels, name="BladeSelector input", - layout=self.layout, allow_dense=self.layout is None or self.layout.dim == self.algebra.dim, ) gate_shape = (1,) * (x.ndim - 2) + (self.channels, self.lane_dim) @@ -110,15 +104,16 @@ def __init__( self.momentum = momentum if not 0.0 <= momentum <= 1.0: raise ValueError(f"momentum must be in [0, 1], got {momentum}") - self.layout = resolve_layer_layout(algebra, layout=layout, grades=grades) - self.lane_dim = lane_dim(algebra, self.layout) + self.storage = resolve_layer_storage(algebra, layout=layout, grades=grades) + self.layout = self.storage.layout + self.lane_dim = self.storage.lane_dim if self.layout is None: self.register_buffer("g0_idx", algebra.grade_indices((0,))) self.register_buffer("g2_idx", algebra.grade_indices((2,))) else: - self.register_buffer("g0_idx", grade_positions(self.layout, 0)) - self.register_buffer("g2_idx", grade_positions(self.layout, 2)) + self.register_buffer("g0_idx", self.storage.grade_positions(0)) + self.register_buffer("g2_idx", self.storage.grade_positions(2)) if self.g0_idx.numel() == 0 or self.g2_idx.numel() == 0: raise ValueError("GeometricNeutralizer layout must include grades 0 and 2") @@ -145,12 +140,10 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: Returns: torch.Tensor: Neutralized multivector. """ - check_multivector_storage( + self.storage.validate_input( x, - self.algebra, channels=self.channels, name="GeometricNeutralizer input", - layout=self.layout, allow_dense=self.layout is None or self.layout.dim == self.algebra.dim, ) diff --git a/layers/primitives/reflection.py b/layers/primitives/reflection.py index 28adac0..52c102b 100644 --- a/layers/primitives/reflection.py +++ b/layers/primitives/reflection.py @@ -11,17 +11,15 @@ from core.foundation.layout import GradeLayout from core.foundation.manifold import MANIFOLD_SPHERE, tag_manifold from core.foundation.module import CliffordModule -from core.planning.action import reflection_vector_matrix +from core.runtime.actions import compact_versor_action from core.runtime.algebra import CliffordAlgebra +from core.runtime.layers import resolve_layer_storage from ._utils import ( cache_matches, - check_multivector_storage, dense_from_indices, grade_indices, - layout_metric_signs, require_positive_int, - resolve_layer_layout, ) @@ -61,12 +59,14 @@ def __init__( """ super().__init__(algebra) self.channels = require_positive_int(channels, "channels") - self.input_layout = resolve_layer_layout(algebra, layout=input_layout, grades=input_grades) - self.output_layout = ( - resolve_layer_layout(algebra, layout=output_layout, grades=output_grades) + self.input_storage = resolve_layer_storage(algebra, layout=input_layout, grades=input_grades) + self.output_storage = ( + resolve_layer_storage(algebra, layout=output_layout, grades=output_grades) if output_layout is not None or output_grades is not None - else self.input_layout + else self.input_storage ) + self.input_layout = self.input_storage.layout + self.output_layout = self.output_storage.layout self.compact_output = bool(compact_output) self.register_buffer("vector_indices", grade_indices(algebra, 1, name="vector grade")) @@ -116,12 +116,10 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: Returns: torch.Tensor: Reflected input [Batch, Channels, Dim]. """ - is_compact = check_multivector_storage( + is_compact = self.input_storage.validate_input( x, - self.algebra, channels=self.channels, name="ReflectionLayer input", - layout=self.input_layout, allow_dense=self.input_layout is None or self.input_layout.dim == self.algebra.dim, ) if is_compact: @@ -155,20 +153,16 @@ def _forward_compact(self, x: torch.Tensor) -> torch.Tensor: """Apply compact reflection through the induced vector action.""" if self.input_layout is None: raise ValueError("ReflectionLayer compact input requires input_layout or input_grades") - output_layout = self.output_layout if self.output_layout is not None else self.input_layout - - normals = self.vector_weights.to(device=x.device, dtype=x.dtype) - signs = layout_metric_signs(self.vector_layout, device=x.device, dtype=x.dtype) - norm_sq = (normals * normals * signs).sum(dim=-1, keepdim=True) - scale = norm_sq.abs().clamp_min(1e-12).sqrt() - normals = normals / scale - matrix = reflection_vector_matrix(normals, vector_layout=self.vector_layout, eps=self.algebra.eps_sq) - return self.algebra.planned_linear_action( + if self.output_layout is None: + raise ValueError("ReflectionLayer compact output requires output_layout or output_grades") + return compact_versor_action( + self.algebra, x, - matrix, + self.vector_weights, + grade=1, input_layout=self.input_layout, - output_layout=output_layout, - input_compact=True, + output_layout=self.output_layout, + parameter_layout=self.vector_layout, compact_output=self.compact_output, ) diff --git a/layers/primitives/rotor.py b/layers/primitives/rotor.py index 931d734..5a29040 100644 --- a/layers/primitives/rotor.py +++ b/layers/primitives/rotor.py @@ -11,17 +11,15 @@ from core.foundation.layout import GradeLayout from core.foundation.manifold import MANIFOLD_SPIN, tag_manifold from core.foundation.module import CliffordModule -from core.planning.action import bivector_vector_generator, reflection_vector_matrix +from core.runtime.actions import compact_versor_action from core.runtime.algebra import CliffordAlgebra +from core.runtime.layers import resolve_layer_storage from ._utils import ( cache_matches, - check_multivector_storage, dense_from_indices, grade_indices, - layout_metric_signs, require_positive_int, - resolve_layer_layout, ) @@ -68,12 +66,14 @@ def __init__( super().__init__(algebra) self.channels = require_positive_int(channels, "channels") self.grade = int(grade) - self.input_layout = resolve_layer_layout(algebra, layout=input_layout, grades=input_grades) - self.output_layout = ( - resolve_layer_layout(algebra, layout=output_layout, grades=output_grades) + self.input_storage = resolve_layer_storage(algebra, layout=input_layout, grades=input_grades) + self.output_storage = ( + resolve_layer_storage(algebra, layout=output_layout, grades=output_grades) if output_layout is not None or output_grades is not None - else self.input_layout + else self.input_storage ) + self.input_layout = self.input_storage.layout + self.output_layout = self.output_storage.layout self.compact_output = bool(compact_output) self.register_buffer("grade_indices", grade_indices(algebra, self.grade)) @@ -148,12 +148,10 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: Returns: torch.Tensor: Transformed input [Batch, Channels, Dim]. """ - is_compact = check_multivector_storage( + is_compact = self.input_storage.validate_input( x, - self.algebra, channels=self.channels, name="RotorLayer input", - layout=self.input_layout, allow_dense=self.input_layout is None or self.input_layout.dim == self.algebra.dim, ) if is_compact: @@ -182,38 +180,18 @@ def _forward_compact(self, x: torch.Tensor) -> torch.Tensor: """Apply a compact grade-preserving versor action.""" if self.input_layout is None: raise ValueError("RotorLayer compact input requires input_layout or input_grades") - output_layout = self.output_layout if self.output_layout is not None else self.input_layout - - if self.grade == 2: - weights = self.grade_weights.to(device=x.device, dtype=x.dtype) - generator = bivector_vector_generator(weights, bivector_layout=self.parameter_layout) - matrix = torch.matrix_exp(generator) - return self.algebra.planned_linear_action( - x, - matrix, - input_layout=self.input_layout, - output_layout=output_layout, - input_compact=True, - compact_output=self.compact_output, - ) - - if self.grade == 1: - weights = self.grade_weights.to(device=x.device, dtype=x.dtype) - signs = layout_metric_signs(self.parameter_layout, device=x.device, dtype=x.dtype) - norm_sq = (weights * weights * signs).sum(dim=-1, keepdim=True) - scale = norm_sq.abs().clamp_min(1e-12).sqrt() - normals = weights / scale - matrix = reflection_vector_matrix(normals, vector_layout=self.parameter_layout, eps=self.algebra.eps_sq) - return self.algebra.planned_linear_action( - x, - matrix, - input_layout=self.input_layout, - output_layout=output_layout, - input_compact=True, - compact_output=self.compact_output, - ) - - raise ValueError("RotorLayer compact execution currently supports grade=1 and grade=2 versors") + if self.output_layout is None: + raise ValueError("RotorLayer compact output requires output_layout or output_grades") + return compact_versor_action( + self.algebra, + x, + self.grade_weights, + grade=self.grade, + input_layout=self.input_layout, + output_layout=self.output_layout, + parameter_layout=self.parameter_layout, + compact_output=self.compact_output, + ) def train(self, mode: bool = True): """Invalidate versor cache when switching to train mode.""" From f3a8b594204ec52967a9d2ddfef338004e73cbec Mon Sep 17 00:00:00 2001 From: Concode0 Date: Tue, 19 May 2026 22:43:22 +0900 Subject: [PATCH 2/4] test: lock pairwise planning cache reuse --- tests/test_framework_pipeline.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/test_framework_pipeline.py b/tests/test_framework_pipeline.py index a566a98..8c95711 100644 --- a/tests/test_framework_pipeline.py +++ b/tests/test_framework_pipeline.py @@ -59,6 +59,8 @@ def test_product_layer_pairwise_compact_widths_match_dense_reference(): ) actual = layer(left, right) + cache_size = len(context.planner._product_executors) + repeated = layer(left, right) expected_dense = dense.wedge( left_layout.dense(left).unsqueeze(2), right_layout.dense(right).unsqueeze(1), @@ -70,6 +72,9 @@ def test_product_layer_pairwise_compact_widths_match_dense_reference(): assert actual.shape == (2, 3, 4, output_layout.dim) assert torch.allclose(actual, expected) + assert torch.allclose(repeated, actual) + assert cache_size == 1 + assert len(context.planner._product_executors) == cache_size def test_compact_layer_pipeline_trains_with_riemannian_optimizer_factory(): From 9dc74a597355f7788a80d7b6d86031ac0dc5646b Mon Sep 17 00:00:00 2001 From: Concode0 Date: Tue, 19 May 2026 23:10:24 +0900 Subject: [PATCH 3/4] refactor: centralize versor layout dispatch in runtime --- core/runtime/__init__.py | 14 +- core/runtime/actions.py | 270 ++++++++++++++++++++++++++++++- core/runtime/layers.py | 9 +- core/runtime/projected.py | 14 ++ layers/primitives/multi_rotor.py | 84 +++------- layers/primitives/reflection.py | 52 ++---- layers/primitives/rotor.py | 65 +++----- tests/test_framework_pipeline.py | 23 +++ tests/test_layers.py | 46 ++++++ tests/test_rotor_gadget.py | 11 ++ 10 files changed, 430 insertions(+), 158 deletions(-) diff --git a/core/runtime/__init__.py b/core/runtime/__init__.py index c2d7dc1..03b46ff 100644 --- a/core/runtime/__init__.py +++ b/core/runtime/__init__.py @@ -8,7 +8,15 @@ """Runtime algebra hosts and dense reference operations.""" from .accessors import as_multivector, compact_values, grade_indices, hermitian_signs, materialize_dense, resolve_layout -from .actions import compact_multi_versor_action, compact_versor_action, versor_vector_matrix +from .actions import ( + apply_multi_versor_action, + apply_versor_action, + compact_multi_versor_action, + compact_versor_action, + dense_versor_factors, + grade_norms, + versor_vector_matrix, +) from .algebra import CliffordAlgebra from .context import AlgebraContext from .layers import LayerStorage, resolve_layer_layout, resolve_layer_storage @@ -22,9 +30,13 @@ "Multivector", "as_multivector", "compact_values", + "apply_multi_versor_action", + "apply_versor_action", "compact_multi_versor_action", "compact_versor_action", + "dense_versor_factors", "grade_indices", + "grade_norms", "hermitian_signs", "materialize_dense", "resolve_layer_layout", diff --git a/core/runtime/actions.py b/core/runtime/actions.py index ba7b8c0..c8a3040 100644 --- a/core/runtime/actions.py +++ b/core/runtime/actions.py @@ -1,14 +1,154 @@ -"""Runtime actions used by compact-capable primitive layers.""" +"""Runtime actions shared by dense and compact algebra hosts.""" from __future__ import annotations import torch from core.foundation.layout import GradeLayout +from core.foundation.validation import check_multivector from core.planning.action import bivector_vector_generator, metric_self_signs, reflection_vector_matrix from core.runtime.accessors import materialize_dense +def apply_versor_action( + algebra, + values: torch.Tensor, + weights: torch.Tensor, + *, + grade: int, + input_grades=None, + output_grades=None, + input_layout: GradeLayout | None = None, + output_layout: GradeLayout | None = None, + parameter_layout: GradeLayout | None = None, + compact_output: bool = False, + channels: int | None = None, + name: str = "versor_action", + dense_cache: tuple[torch.Tensor, torch.Tensor] | None = None, + cache_dense: bool = False, + return_cache: bool = False, +) -> torch.Tensor | tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]: + """Apply one versor action while the algebra host chooses storage execution.""" + input_layout = _declared_layout(algebra, input_grades, input_layout) + output_layout = _declared_layout(algebra, output_grades, output_layout) or input_layout + parameter_layout = parameter_layout or algebra.layout((int(grade),)) + + input_compact = _validate_action_values( + algebra, + values, + layout=input_layout, + channels=channels, + name=name, + ) + if input_compact: + output = compact_versor_action( + algebra, + values, + weights, + grade=grade, + input_layout=input_layout, + output_layout=output_layout, + parameter_layout=parameter_layout, + compact_output=compact_output, + ) + return (output, dense_cache) if return_cache else output + + _require_dense_action(algebra, name, multi=False) + left, right, next_cache = _dense_versor_factors( + algebra, + values, + weights, + grade=grade, + parameter_layout=parameter_layout, + dense_cache=dense_cache, + cache_dense=cache_dense, + ) + output = algebra.per_channel_sandwich(left, values, right) + output = _project_dense_action_output(algebra, output, output_layout=output_layout, compact_output=compact_output) + return (output, next_cache) if return_cache else output + + +def apply_multi_versor_action( + algebra, + values: torch.Tensor, + weights: torch.Tensor, + mix: torch.Tensor, + *, + grade: int, + input_grades=None, + output_grades=None, + input_layout: GradeLayout | None = None, + output_layout: GradeLayout | None = None, + parameter_layout: GradeLayout | None = None, + compact_output: bool = False, + channels: int | None = None, + name: str = "multi_versor_action", + dense_cache: tuple[torch.Tensor, torch.Tensor] | None = None, + cache_dense: bool = False, + return_cache: bool = False, +) -> torch.Tensor | tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]: + """Apply a weighted versor superposition with host-owned storage dispatch.""" + input_layout = _declared_layout(algebra, input_grades, input_layout) + output_layout = _declared_layout(algebra, output_grades, output_layout) or input_layout + parameter_layout = parameter_layout or algebra.layout((int(grade),)) + + input_compact = _validate_action_values( + algebra, + values, + layout=input_layout, + channels=channels, + name=name, + ) + if input_compact: + output = compact_multi_versor_action( + algebra, + values, + weights, + mix, + grade=grade, + input_layout=input_layout, + output_layout=output_layout, + parameter_layout=parameter_layout, + compact_output=compact_output, + ) + return (output, dense_cache) if return_cache else output + + _require_dense_action(algebra, name, multi=True) + left, right, next_cache = _dense_versor_factors( + algebra, + values, + weights, + grade=grade, + parameter_layout=parameter_layout, + dense_cache=dense_cache, + cache_dense=cache_dense, + ) + versored = algebra.multi_rotor_sandwich(left, values, right) + output = torch.einsum("ck,...cke->...ce", mix.to(device=values.device, dtype=values.dtype), versored) + output = _project_dense_action_output(algebra, output, output_layout=output_layout, compact_output=compact_output) + return (output, next_cache) if return_cache else output + + +def grade_norms( + algebra, + values: torch.Tensor, + *, + input_grades=None, + layout: GradeLayout | None = None, +) -> torch.Tensor: + """Return per-grade coefficient norms for dense or compact values.""" + layout = _declared_layout(algebra, input_grades, layout) + input_compact = _values_are_compact(algebra, values, layout) + if input_compact: + return compact_grade_norms(algebra, values, layout) + if hasattr(algebra, "get_grade_norms"): + return algebra.get_grade_norms(values) + + check_multivector(values, algebra, "grade_norms(values)") + full_layout = algebra.layout(range(algebra.num_grades)) + return compact_grade_norms(algebra, values, full_layout) + + def compact_versor_action( algebra, values: torch.Tensor, @@ -94,3 +234,131 @@ def versor_vector_matrix(algebra, weights: torch.Tensor, *, grade: int, paramete def parameter_layout_signs(layout: GradeLayout, *, device=None, dtype=None) -> torch.Tensor: """Return basis self-product signs for compact parameter weights.""" return metric_self_signs(layout, device=device, dtype=dtype) + + +def compact_grade_norms(algebra, values: torch.Tensor, layout: GradeLayout) -> torch.Tensor: + """Return per-grade coefficient norms for compact values.""" + flat = values.pow(2).reshape(-1, layout.dim) + grade_ids = layout.grade_indices_tensor(device=values.device).unsqueeze(0).expand_as(flat) + result = values.new_zeros(flat.shape[0], algebra.num_grades) + result.scatter_add_(1, grade_ids, flat) + return result.reshape(*values.shape[:-1], algebra.num_grades).clamp(min=algebra.eps).sqrt() + + +def dense_versor_factors( + algebra, + weights: torch.Tensor, + *, + grade: int, + parameter_layout: GradeLayout, +) -> tuple[torch.Tensor, torch.Tensor]: + """Return dense left/right factors for a parameterized versor.""" + versor = parameter_layout.dense(weights) + grade = int(grade) + + if grade == 2: + rotor = algebra.exp(-0.5 * versor) + return rotor, algebra.reverse(rotor) + + if grade == 1: + norm_sq = algebra.norm_sq(versor) + scale = norm_sq.abs().clamp(min=1e-12).sqrt() + versor = versor / scale + else: + norm = versor.norm(dim=-1, keepdim=True).clamp(min=1e-8) + versor = versor / norm + return algebra.grade_involution(versor), algebra.blade_inverse(versor) + + +def _dense_versor_factors( + algebra, + values: torch.Tensor, + weights: torch.Tensor, + *, + grade: int, + parameter_layout: GradeLayout, + dense_cache: tuple[torch.Tensor, torch.Tensor] | None, + cache_dense: bool, +) -> tuple[torch.Tensor, torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]: + if _cache_matches(dense_cache, values): + left, right = dense_cache + else: + dense_weights = weights.to(device=values.device, dtype=values.dtype) + left, right = dense_versor_factors( + algebra, + dense_weights, + grade=grade, + parameter_layout=parameter_layout, + ) + return left, right, (left, right) if cache_dense else None + + +def _declared_layout(algebra, grades, layout: GradeLayout | None) -> GradeLayout | None: + if hasattr(algebra, "_declared_layout"): + return algebra._declared_layout(grades, layout) + if layout is not None: + return layout + if grades is not None: + return algebra.layout(grades) + default_grades = getattr(algebra, "_default_grades", None) + if default_grades is None: + return None + return algebra.layout(default_grades) + + +def _validate_action_values( + algebra, + values: torch.Tensor, + *, + layout: GradeLayout | None, + channels: int | None, + name: str, +) -> bool: + if values.ndim < 3: + raise ValueError(f"{name}: expected ndim >= 3, got shape {tuple(values.shape)}") + if channels is not None and values.shape[-2] != channels: + raise ValueError(f"{name}: expected {channels} channels, got {values.shape[-2]} (shape {tuple(values.shape)})") + + if _values_are_compact(algebra, values, layout): + return True + if values.shape[-1] == algebra.dim: + return False + + expected = [str(algebra.dim)] + if layout is not None: + expected.insert(0, f"{layout.dim} for grades {layout.grades}") + raise ValueError(f"{name}: last dim must be {' or '.join(expected)}, got {values.shape[-1]}") + + +def _values_are_compact(algebra, values: torch.Tensor, layout: GradeLayout | None) -> bool: + if layout is None or values.shape[-1] != layout.dim: + return False + return layout.dim != algebra.dim or not hasattr(algebra, "per_channel_sandwich") + + +def _require_dense_action(algebra, name: str, *, multi: bool) -> None: + if not hasattr(algebra, "per_channel_sandwich"): + raise ValueError(f"{name}: dense execution requires CliffordAlgebra; declare grades/layout for compact use") + if multi and not hasattr(algebra, "multi_rotor_sandwich"): + raise ValueError(f"{name}: dense execution requires CliffordAlgebra action kernels") + + +def _project_dense_action_output( + algebra, + output: torch.Tensor, + *, + output_layout: GradeLayout | None, + compact_output: bool, +) -> torch.Tensor: + if output_layout is None: + return output + compact = output_layout.compact(output) + if compact_output: + return compact + return materialize_dense(algebra, compact, layout=output_layout) + + +def _cache_matches(cache: tuple[torch.Tensor, ...] | None, reference: torch.Tensor) -> bool: + if cache is None: + return False + return all(tensor.device == reference.device and tensor.dtype == reference.dtype for tensor in cache) diff --git a/core/runtime/layers.py b/core/runtime/layers.py index 2bfa4ad..2b01fcc 100644 --- a/core/runtime/layers.py +++ b/core/runtime/layers.py @@ -8,6 +8,7 @@ from core.foundation.layout import GradeLayout from core.planning.action import metric_self_signs +from core.runtime.actions import compact_grade_norms @dataclass(frozen=True) @@ -86,12 +87,8 @@ def grade_positions(self, grade: int, *, device=None) -> torch.Tensor: def compact_grade_norms(self, values: torch.Tensor) -> torch.Tensor: """Return per-grade coefficient norms for compact values.""" if self.layout is None: - return self.algebra.get_grade_norms(values) - flat = values.pow(2).reshape(-1, self.layout.dim) - grade_ids = self.layout.grade_indices_tensor(device=values.device).unsqueeze(0).expand_as(flat) - result = values.new_zeros(flat.shape[0], self.algebra.num_grades) - result.scatter_add_(1, grade_ids, flat) - return result.reshape(*values.shape[:-1], self.algebra.num_grades).clamp(min=self.algebra.eps).sqrt() + return self.algebra.grade_norms(values) + return compact_grade_norms(self.algebra, values, self.layout) def metric_signs(self, *, device=None, dtype=None) -> torch.Tensor: """Return basis self-product signs for this storage.""" diff --git a/core/runtime/projected.py b/core/runtime/projected.py index d8c30a3..f79aeb4 100644 --- a/core/runtime/projected.py +++ b/core/runtime/projected.py @@ -21,6 +21,8 @@ from core.runtime.accessors import hermitian_signs as _hermitian_signs from core.runtime.accessors import materialize_dense from core.runtime.accessors import resolve_layout as _resolve_layout +from core.runtime.actions import apply_multi_versor_action, apply_versor_action +from core.runtime.actions import grade_norms as _grade_norms class AlgebraRuntimeMixin: @@ -96,6 +98,18 @@ def hermitian_signs( """Return Hermitian signs for a dense or compact layout.""" return _hermitian_signs(self, layout=layout, grades=grades, device=device, dtype=dtype) + def versor_action(self, values: torch.Tensor, weights: torch.Tensor, **kwargs): + """Apply a parameterized versor action through the host storage dispatcher.""" + return apply_versor_action(self, values, weights, **kwargs) + + def multi_versor_action(self, values: torch.Tensor, weights: torch.Tensor, mix: torch.Tensor, **kwargs): + """Apply a weighted versor superposition through the host storage dispatcher.""" + return apply_multi_versor_action(self, values, weights, mix, **kwargs) + + def grade_norms(self, values: torch.Tensor, **kwargs) -> torch.Tensor: + """Return per-grade norms for dense or compact values.""" + return _grade_norms(self, values, **kwargs) + def projected_product( self, A: torch.Tensor, diff --git a/layers/primitives/multi_rotor.py b/layers/primitives/multi_rotor.py index 38383ed..608e8f5 100644 --- a/layers/primitives/multi_rotor.py +++ b/layers/primitives/multi_rotor.py @@ -16,13 +16,11 @@ from core.foundation.layout import GradeLayout from core.foundation.manifold import MANIFOLD_SPIN, tag_manifold from core.foundation.module import CliffordModule -from core.runtime.actions import compact_multi_versor_action +from core.runtime.actions import dense_versor_factors from core.runtime.algebra import CliffordAlgebra from core.runtime.layers import resolve_layer_storage from ._utils import ( - cache_matches, - dense_from_indices, grade_indices, require_positive_int, ) @@ -130,15 +128,12 @@ def _compute_versors(self, device, dtype): Tuple[Tensor, Tensor]: (V_left [K, dim], V_right [K, dim]) """ weights = self.rotor_grade_weights.to(device=device, dtype=dtype) - V = dense_from_indices(weights, self.grade_indices, self.algebra.dim) - - if self.grade == 2: - R = self.algebra.exp(-0.5 * V) # [K, D] - return R, self.algebra.reverse(R) - else: - norm = V.norm(dim=-1, keepdim=True).clamp(min=1e-8) - V = V / norm - return self.algebra.grade_involution(V), self.algebra.blade_inverse(V) + return dense_versor_factors( + self.algebra, + weights, + grade=self.grade, + parameter_layout=self.parameter_layout, + ) def forward(self, x: torch.Tensor, return_invariants: bool = False) -> torch.Tensor: """Apply weighted multi-versor superposition. @@ -152,61 +147,12 @@ def forward(self, x: torch.Tensor, return_invariants: bool = False) -> torch.Ten Returns: torch.Tensor: Transformed output [Batch, Channels, Dim]. """ - is_compact = self.input_storage.validate_input( - x, - channels=self.channels, - name="MultiRotorLayer input", - allow_dense=self.input_layout is None or self.input_layout.dim == self.algebra.dim, - ) - if is_compact: - out = self._forward_compact(x) - if return_invariants: - if not self.compact_output: - return self.algebra.get_grade_norms(out) - return self.output_storage.compact_grade_norms(out) - return out - if not hasattr(self.algebra, "multi_rotor_sandwich"): - raise ValueError( - "MultiRotorLayer dense execution requires CliffordAlgebra; declare input_grades for compact use." - ) - cache = ( (self._cached_V_left, self._cached_V_right) - if self._cached_V_left is not None and self._cached_V_right is not None + if not self.training and self._cached_V_left is not None and self._cached_V_right is not None else None ) - if not self.training and cache_matches(cache, x): - V_left, V_right = self._cached_V_left, self._cached_V_right - else: - V_left, V_right = self._compute_versors(x.device, x.dtype) - if not self.training: - self._cached_V_left = V_left - self._cached_V_right = V_right - - # Action-matrix sandwich: build K matrices once, apply via einsum - versored_x = self.algebra.multi_rotor_sandwich( - V_left, - x, - V_right, - ) # [B, C, K, D] - - # Weighted superposition - weights = self.weights.to(device=x.device, dtype=x.dtype) - out = torch.einsum("ck,...cke->...ce", weights, versored_x) - - if return_invariants: - return self.algebra.get_grade_norms(out) - - return out - - def _forward_compact(self, x: torch.Tensor) -> torch.Tensor: - """Apply compact weighted superposition of induced versor actions.""" - if self.input_layout is None: - raise ValueError("MultiRotorLayer compact input requires input_layout or input_grades") - if self.output_layout is None: - raise ValueError("MultiRotorLayer compact output requires output_layout or output_grades") - return compact_multi_versor_action( - self.algebra, + out, next_cache = self.algebra.multi_versor_action( x, self.rotor_grade_weights, self.weights, @@ -215,7 +161,19 @@ def _forward_compact(self, x: torch.Tensor) -> torch.Tensor: output_layout=self.output_layout, parameter_layout=self.parameter_layout, compact_output=self.compact_output, + channels=self.channels, + name="MultiRotorLayer input", + dense_cache=cache, + cache_dense=not self.training, + return_cache=True, ) + if not self.training and next_cache is not None: + self._cached_V_left, self._cached_V_right = next_cache + + if return_invariants: + return self.algebra.grade_norms(out, layout=self.output_layout) + + return out def train(self, mode: bool = True): """Invalidate versor cache when switching to train mode.""" diff --git a/layers/primitives/reflection.py b/layers/primitives/reflection.py index 52c102b..5654f70 100644 --- a/layers/primitives/reflection.py +++ b/layers/primitives/reflection.py @@ -11,13 +11,10 @@ from core.foundation.layout import GradeLayout from core.foundation.manifold import MANIFOLD_SPHERE, tag_manifold from core.foundation.module import CliffordModule -from core.runtime.actions import compact_versor_action from core.runtime.algebra import CliffordAlgebra from core.runtime.layers import resolve_layer_storage from ._utils import ( - cache_matches, - dense_from_indices, grade_indices, require_positive_int, ) @@ -97,7 +94,7 @@ def _build_vectors(self, device, dtype): Tuple of (n, n_inv) each [C, dim]. """ weights = self.vector_weights.to(device=device, dtype=dtype) - n = dense_from_indices(weights, self.vector_indices, self.algebra.dim) + n = self.vector_layout.dense(weights) # Normalize: n_hat = n / sqrt(|_0|) n_sq = self.algebra.norm_sq(n) # [C, 1] @@ -116,47 +113,12 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: Returns: torch.Tensor: Reflected input [Batch, Channels, Dim]. """ - is_compact = self.input_storage.validate_input( - x, - channels=self.channels, - name="ReflectionLayer input", - allow_dense=self.input_layout is None or self.input_layout.dim == self.algebra.dim, - ) - if is_compact: - return self._forward_compact(x) - if not hasattr(self.algebra, "per_channel_sandwich"): - raise ValueError( - "ReflectionLayer dense execution requires CliffordAlgebra; declare input_grades for compact use." - ) - cache = ( (self._cached_n, self._cached_n_inv) - if self._cached_n is not None and self._cached_n_inv is not None + if not self.training and self._cached_n is not None and self._cached_n_inv is not None else None ) - if not self.training and cache_matches(cache, x): - n, n_inv = self._cached_n, self._cached_n_inv - else: - n, n_inv = self._build_vectors(x.device, x.dtype) - if not self.training: - self._cached_n = n - self._cached_n_inv = n_inv - - # grade_involution(n) = -n for grade-1 vectors - n_hat = -n # [C, dim] - - # Per-channel reflection via two GPs: (-n) * x * n^{-1} - # Use per_channel_sandwich with n_hat as "R" and n_inv as "R_rev" - return self.algebra.per_channel_sandwich(n_hat, x, n_inv) - - def _forward_compact(self, x: torch.Tensor) -> torch.Tensor: - """Apply compact reflection through the induced vector action.""" - if self.input_layout is None: - raise ValueError("ReflectionLayer compact input requires input_layout or input_grades") - if self.output_layout is None: - raise ValueError("ReflectionLayer compact output requires output_layout or output_grades") - return compact_versor_action( - self.algebra, + out, next_cache = self.algebra.versor_action( x, self.vector_weights, grade=1, @@ -164,7 +126,15 @@ def _forward_compact(self, x: torch.Tensor) -> torch.Tensor: output_layout=self.output_layout, parameter_layout=self.vector_layout, compact_output=self.compact_output, + channels=self.channels, + name="ReflectionLayer input", + dense_cache=cache, + cache_dense=not self.training, + return_cache=True, ) + if not self.training and next_cache is not None: + self._cached_n, self._cached_n_inv = next_cache + return out def train(self, mode: bool = True): """Override to invalidate cache when switching to train mode.""" diff --git a/layers/primitives/rotor.py b/layers/primitives/rotor.py index 5a29040..230021d 100644 --- a/layers/primitives/rotor.py +++ b/layers/primitives/rotor.py @@ -11,13 +11,11 @@ from core.foundation.layout import GradeLayout from core.foundation.manifold import MANIFOLD_SPIN, tag_manifold from core.foundation.module import CliffordModule -from core.runtime.actions import compact_versor_action +from core.runtime.actions import dense_versor_factors from core.runtime.algebra import CliffordAlgebra from core.runtime.layers import resolve_layer_storage from ._utils import ( - cache_matches, - dense_from_indices, grade_indices, require_positive_int, ) @@ -113,7 +111,7 @@ def reset_parameters(self): def _build_grade_element(self, device, dtype): """Scatter grade_weights into full multivector dimension [channels, dim].""" weights = self.grade_weights.to(device=device, dtype=dtype) - return dense_from_indices(weights, self.grade_indices, self.algebra.dim) + return self.parameter_layout.dense(weights) def _compute_versors(self, device, dtype): """Compute left and right factors for per_channel_sandwich. @@ -126,16 +124,13 @@ def _compute_versors(self, device, dtype): Returns: Tuple[Tensor, Tensor]: (V_left [C, dim], V_right [C, dim]) """ - V = self._build_grade_element(device, dtype) - if self.grade == 2: - R = self.algebra.exp(-0.5 * V) - return R, self.algebra.reverse(R) - else: - # Normalize per channel so blade_inverse is exact. - # For a unit-norm grade-k element, V * V_rev = scalar everywhere. - norm = V.norm(dim=-1, keepdim=True).clamp(min=1e-8) - V = V / norm - return self.algebra.grade_involution(V), self.algebra.blade_inverse(V) + weights = self.grade_weights.to(device=device, dtype=dtype) + return dense_versor_factors( + self.algebra, + weights, + grade=self.grade, + parameter_layout=self.parameter_layout, + ) def forward(self, x: torch.Tensor) -> torch.Tensor: """Apply versor product x' = hat(V) x V^{-1} (= RxR~ for grade=2). @@ -148,42 +143,12 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: Returns: torch.Tensor: Transformed input [Batch, Channels, Dim]. """ - is_compact = self.input_storage.validate_input( - x, - channels=self.channels, - name="RotorLayer input", - allow_dense=self.input_layout is None or self.input_layout.dim == self.algebra.dim, - ) - if is_compact: - return self._forward_compact(x) - if not hasattr(self.algebra, "per_channel_sandwich"): - raise ValueError( - "RotorLayer dense execution requires CliffordAlgebra; declare input_grades for compact use." - ) - cache = ( (self._cached_V_left, self._cached_V_right) - if self._cached_V_left is not None and self._cached_V_right is not None + if not self.training and self._cached_V_left is not None and self._cached_V_right is not None else None ) - if not self.training and cache_matches(cache, x): - V_left, V_right = self._cached_V_left, self._cached_V_right - else: - V_left, V_right = self._compute_versors(x.device, x.dtype) - if not self.training: - self._cached_V_left = V_left - self._cached_V_right = V_right - - return self.algebra.per_channel_sandwich(V_left, x, V_right) - - def _forward_compact(self, x: torch.Tensor) -> torch.Tensor: - """Apply a compact grade-preserving versor action.""" - if self.input_layout is None: - raise ValueError("RotorLayer compact input requires input_layout or input_grades") - if self.output_layout is None: - raise ValueError("RotorLayer compact output requires output_layout or output_grades") - return compact_versor_action( - self.algebra, + out, next_cache = self.algebra.versor_action( x, self.grade_weights, grade=self.grade, @@ -191,7 +156,15 @@ def _forward_compact(self, x: torch.Tensor) -> torch.Tensor: output_layout=self.output_layout, parameter_layout=self.parameter_layout, compact_output=self.compact_output, + channels=self.channels, + name="RotorLayer input", + dense_cache=cache, + cache_dense=not self.training, + return_cache=True, ) + if not self.training and next_cache is not None: + self._cached_V_left, self._cached_V_right = next_cache + return out def train(self, mode: bool = True): """Invalidate versor cache when switching to train mode.""" diff --git a/tests/test_framework_pipeline.py b/tests/test_framework_pipeline.py index 8c95711..0a52d25 100644 --- a/tests/test_framework_pipeline.py +++ b/tests/test_framework_pipeline.py @@ -6,6 +6,7 @@ from core.runtime.algebra import CliffordAlgebra from core.runtime.context import AlgebraContext from layers import ProductLayer, WedgeLayer +from layers.blocks.multi_rotor_ffn import MultiRotorFFN from optimizers import make_riemannian_optimizer pytestmark = pytest.mark.unit @@ -125,6 +126,28 @@ def forward(self, left_vector, right_vector, third_vector): assert torch.isfinite(model.scale).all() +def test_rotor_backend_block_trains_with_riemannian_optimizer_factory(): + algebra = CliffordAlgebra(p=3, q=0, device="cpu") + model = MultiRotorFFN( + algebra, + channels=2, + ffn_mult=2, + num_rotors=2, + use_rotor_backend=True, + ) + optimizer = make_riemannian_optimizer(model, algebra, optimizer="adam", lr=0.01) + x = torch.randn(4, 2, algebra.dim) + + output = model(x) + loss = output.square().mean() + loss.backward() + optimizer.step() + + assert output.shape == x.shape + assert any(group.get("manifold") == "spin" for group in optimizer.param_groups) + assert all(torch.isfinite(parameter).all() for parameter in model.parameters()) + + def test_product_layer_uses_context_planning_limits(): limits = PlanningLimits(warn_lanes=32, max_lanes=512, warn_pairs=32, max_pairs=64) context = AlgebraContext(p=16, q=0, device="cpu", planning_limits=limits) diff --git a/tests/test_layers.py b/tests/test_layers.py index 97e86f8..2d83d94 100644 --- a/tests/test_layers.py +++ b/tests/test_layers.py @@ -140,6 +140,21 @@ def test_compact_rotor_multigrade_matches_dense_reference(self): assert actual.shape == x.shape assert torch.allclose(actual, expected, atol=1e-4) + def test_declared_rotor_layout_accepts_dense_input_and_returns_compact(self): + dense = CliffordAlgebra(3, 0, device="cpu") + layout = dense.layout((1,)) + x = dense.embed_vector(torch.randn(2, 4, dense.n)) + + declared_layer = RotorLayer(dense, 4, input_layout=layout) + reference_layer = RotorLayer(dense, 4) + reference_layer.grade_weights.data.copy_(declared_layer.grade_weights.data) + + actual = declared_layer(x) + expected = layout.compact(reference_layer(x)) + + assert actual.shape == (2, 4, layout.dim) + assert torch.allclose(actual, expected, atol=1e-4) + def test_multi_rotor_shape(self, algebra_3d): x = torch.randn(4, 5, 8) layer = MultiRotorLayer(algebra_3d, 5, num_rotors=4) @@ -333,6 +348,21 @@ def test_compact_reflection_matches_dense_reference(self): assert actual.shape == x.shape assert torch.allclose(actual, expected, atol=1e-4) + def test_declared_reflection_layout_accepts_dense_input_and_returns_compact(self): + dense = CliffordAlgebra(3, 0, device="cpu") + layout = dense.layout((1,)) + x = dense.embed_vector(torch.randn(2, 4, dense.n)) + + declared_layer = ReflectionLayer(dense, channels=4, input_layout=layout) + reference_layer = ReflectionLayer(dense, channels=4) + reference_layer.vector_weights.data.copy_(declared_layer.vector_weights.data) + + actual = declared_layer(x) + expected = layout.compact(reference_layer(x)) + + assert actual.shape == (2, 4, layout.dim) + assert torch.allclose(actual, expected, atol=1e-4) + def test_reflection_preserves_norm(self, algebra_3d): C = 3 layer = ReflectionLayer(algebra_3d, channels=C) @@ -444,6 +474,22 @@ def test_compile_multi_rotor_layer(self, algebra_3d): y = compiled(x) assert y.shape == (2, 4, 8) + def test_compile_compact_versor_layers(self): + """Compact versor layers compile through the polymorphic core dispatcher.""" + context = AlgebraContext(3, 0, device="cpu") + layout = context.layout((1,)) + x = torch.randn(2, 4, layout.dim) + + layers = ( + RotorLayer(context, channels=4, input_layout=layout), + ReflectionLayer(context, channels=4, input_layout=layout), + MultiRotorLayer(context, channels=4, num_rotors=2, input_layout=layout), + ) + for layer in layers: + compiled = torch.compile(layer, backend="aot_eager", fullgraph=True) + y = compiled(x) + assert y.shape == x.shape + def test_compile_backward(self, algebra_3d): """Gradients flow through compiled RotorLayer.""" layer = RotorLayer(algebra_3d, channels=4) diff --git a/tests/test_rotor_gadget.py b/tests/test_rotor_gadget.py index 2d5cc23..8c07b94 100644 --- a/tests/test_rotor_gadget.py +++ b/tests/test_rotor_gadget.py @@ -38,6 +38,17 @@ def test_basic_forward(self, algebra_2d): assert out.shape == (batch_size, 8, algebra_2d.dim) + @pytest.mark.skipif(not hasattr(torch, "compile"), reason="torch.compile not available") + def test_fullgraph_compile(self, algebra_3d): + """RotorGadget remains graph-capturable on its dense execution path.""" + layer = RotorGadget(algebra=algebra_3d, in_channels=4, out_channels=4, num_rotor_pairs=2) + compiled = torch.compile(layer, backend="aot_eager", fullgraph=True) + x = torch.randn(2, 4, algebra_3d.dim) + + y = compiled(x) + + assert y.shape == (2, 4, algebra_3d.dim) + @pytest.mark.parametrize( "in_ch,out_ch,num_pairs", [ From 2b2527b6b76d105cb7610b610fc7e7bcf8e8d559 Mon Sep 17 00:00:00 2001 From: Concode0 Date: Tue, 19 May 2026 23:24:21 +0900 Subject: [PATCH 4/4] perf: vectorize compact multi-versor actions --- core/planning/__init__.py | 8 ++++- core/planning/action.py | 30 ++++++++++++++++ core/runtime/actions.py | 72 +++++++++++++++++++++++++++------------ tests/test_grade_plan.py | 25 ++++++++++++++ 4 files changed, 112 insertions(+), 23 deletions(-) diff --git a/core/planning/__init__.py b/core/planning/__init__.py index cb681bc..d67a976 100644 --- a/core/planning/__init__.py +++ b/core/planning/__init__.py @@ -7,7 +7,12 @@ """Static grade planning and Torch executor lowering.""" -from .action import apply_graded_linear_action, bivector_vector_generator, reflection_vector_matrix +from .action import ( + apply_graded_linear_action, + apply_multi_graded_linear_action, + bivector_vector_generator, + reflection_vector_matrix, +) from .flow import GradeFlow from .layouts import ProductRequest, build_product_request from .planner import GradePlanner @@ -27,6 +32,7 @@ "PlanCost", "DEFAULT_PLANNING_LIMITS", "apply_graded_linear_action", + "apply_multi_graded_linear_action", "bivector_vector_generator", "GradeUnaryExecutor", "GradeUnaryOp", diff --git a/core/planning/action.py b/core/planning/action.py index fe00bb7..6a39568 100644 --- a/core/planning/action.py +++ b/core/planning/action.py @@ -41,6 +41,36 @@ def apply_graded_linear_action( return torch.einsum("coi,...ci->...co", coefficients, values) +def apply_multi_graded_linear_action( + values: torch.Tensor, + matrices: torch.Tensor, + *, + input_layout: GradeLayout, + output_layout: GradeLayout, +) -> torch.Tensor: + """Apply multiple outermorphisms to compact grade lanes. + + ``matrices`` stores ``[actions, n, n]`` vector-space maps. ``values`` + stores ``[..., channels, input_layout.dim]`` compact lanes. The result is + ``[..., channels, actions, output_layout.dim]``. + """ + if input_layout.spec != output_layout.spec: + raise ValueError(f"layout mismatch: {input_layout.spec} vs {output_layout.spec}") + if values.shape[-1] != input_layout.dim: + raise ValueError(f"input compact dimension must be {input_layout.dim}, got {values.shape[-1]}") + if values.ndim < 2: + raise ValueError(f"values must include channel and lane axes, got shape {tuple(values.shape)}") + + spec = input_layout.spec + if matrices.shape[-2:] != (spec.n, spec.n): + raise ValueError(f"matrices trailing shape must be {(spec.n, spec.n)}, got {tuple(matrices.shape[-2:])}") + if matrices.ndim != 3: + raise ValueError(f"matrices must have shape [actions, n, n], got {tuple(matrices.shape)}") + + coefficients = _graded_action_coefficients(matrices, input_layout=input_layout, output_layout=output_layout) + return torch.einsum("koi,...ci->...cko", coefficients, values) + + def bivector_vector_generator(bivectors: torch.Tensor, *, bivector_layout: GradeLayout) -> torch.Tensor: """Return the vector-space generator induced by compact bivectors.""" if bivector_layout.grades != (2,): diff --git a/core/runtime/actions.py b/core/runtime/actions.py index c8a3040..08619f8 100644 --- a/core/runtime/actions.py +++ b/core/runtime/actions.py @@ -6,7 +6,12 @@ from core.foundation.layout import GradeLayout from core.foundation.validation import check_multivector -from core.planning.action import bivector_vector_generator, metric_self_signs, reflection_vector_matrix +from core.planning.action import ( + apply_multi_graded_linear_action, + bivector_vector_generator, + metric_self_signs, + reflection_vector_matrix, +) from core.runtime.accessors import materialize_dense @@ -29,9 +34,15 @@ def apply_versor_action( return_cache: bool = False, ) -> torch.Tensor | tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]: """Apply one versor action while the algebra host chooses storage execution.""" - input_layout = _declared_layout(algebra, input_grades, input_layout) - output_layout = _declared_layout(algebra, output_grades, output_layout) or input_layout - parameter_layout = parameter_layout or algebra.layout((int(grade),)) + input_layout, output_layout, parameter_layout = _action_layouts( + algebra, + grade=grade, + input_grades=input_grades, + output_grades=output_grades, + input_layout=input_layout, + output_layout=output_layout, + parameter_layout=parameter_layout, + ) input_compact = _validate_action_values( algebra, @@ -88,9 +99,15 @@ def apply_multi_versor_action( return_cache: bool = False, ) -> torch.Tensor | tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]: """Apply a weighted versor superposition with host-owned storage dispatch.""" - input_layout = _declared_layout(algebra, input_grades, input_layout) - output_layout = _declared_layout(algebra, output_grades, output_layout) or input_layout - parameter_layout = parameter_layout or algebra.layout((int(grade),)) + input_layout, output_layout, parameter_layout = _action_layouts( + algebra, + grade=grade, + input_grades=input_grades, + output_grades=output_grades, + input_layout=input_layout, + output_layout=output_layout, + parameter_layout=parameter_layout, + ) input_compact = _validate_action_values( algebra, @@ -196,22 +213,17 @@ def compact_multi_versor_action( grade=grade, parameter_layout=parameter_layout, ) - outputs = [] - for index in range(matrices.shape[0]): - matrix = matrices[index].unsqueeze(0).expand(values.shape[-2], -1, -1) - outputs.append( - algebra.planned_linear_action( - values, - matrix, - input_layout=input_layout, - output_layout=output_layout, - input_compact=True, - compact_output=True, - ) - ) + mix = mix.to(device=values.device, dtype=values.dtype) + if mix.shape != (values.shape[-2], matrices.shape[0]): + raise ValueError(f"mix shape must be {(values.shape[-2], matrices.shape[0])}, got {tuple(mix.shape)}") - stacked = torch.stack(outputs, dim=-2) - result = torch.einsum("ck,...ckd->...cd", mix.to(device=values.device, dtype=values.dtype), stacked) + transformed = apply_multi_graded_linear_action( + values, + matrices, + input_layout=input_layout, + output_layout=output_layout, + ) + result = torch.einsum("ck,...cko->...co", mix, transformed) if compact_output: return result return materialize_dense(algebra, result, layout=output_layout) @@ -306,6 +318,22 @@ def _declared_layout(algebra, grades, layout: GradeLayout | None) -> GradeLayout return algebra.layout(default_grades) +def _action_layouts( + algebra, + *, + grade: int, + input_grades, + output_grades, + input_layout: GradeLayout | None, + output_layout: GradeLayout | None, + parameter_layout: GradeLayout | None, +) -> tuple[GradeLayout | None, GradeLayout | None, GradeLayout]: + input_layout = _declared_layout(algebra, input_grades, input_layout) + output_layout = _declared_layout(algebra, output_grades, output_layout) or input_layout + parameter_layout = parameter_layout or algebra.layout((int(grade),)) + return input_layout, output_layout, parameter_layout + + def _validate_action_values( algebra, values: torch.Tensor, diff --git a/tests/test_grade_plan.py b/tests/test_grade_plan.py index cd4d801..a2e16a4 100644 --- a/tests/test_grade_plan.py +++ b/tests/test_grade_plan.py @@ -12,6 +12,7 @@ product_output_grades, ) from core.foundation.layout import AlgebraSpec +from core.planning.action import apply_graded_linear_action, apply_multi_graded_linear_action from core.planning.flow import GradeFlow from core.planning.layouts import build_product_request from core.planning.planner import GradePlanner @@ -353,6 +354,30 @@ def test_product_executor_compact_forward_supports_different_layout_widths(): assert torch.allclose(compact, dense, atol=1e-12, rtol=1e-12) +def test_multi_graded_linear_action_matches_stacked_single_actions(): + algebra = CliffordAlgebra(4, 0, 0, device=DEVICE, dtype=torch.float64) + layout = algebra.layout((0, 1, 2)) + values = torch.randn(2, 3, layout.dim, dtype=torch.float64) + matrices = torch.randn(5, algebra.n, algebra.n, dtype=torch.float64) + + actual = apply_multi_graded_linear_action(values, matrices, input_layout=layout, output_layout=layout) + expected = torch.stack( + [ + apply_graded_linear_action( + values, + matrix.unsqueeze(0).expand(values.shape[-2], -1, -1), + input_layout=layout, + output_layout=layout, + ) + for matrix in matrices + ], + dim=-2, + ) + + assert actual.shape == (2, 3, 5, layout.dim) + assert torch.allclose(actual, expected, atol=1e-12, rtol=1e-12) + + def test_algebra_projected_product_matches_dense_kernel_and_compact_output(): algebra = CliffordAlgebra(4, 1, 1, device=DEVICE, dtype=torch.float64) A = _grade_only_input(algebra, 2, (1,), seed=113)