morphis 0.7.0__py3-none-any.whl → 0.9.0__py3-none-any.whl

morphis/algebra/__init__.py

@@ -2,9 +2,13 @@
 Linear Algebra Module
 
 Provides structured linear algebra utilities for geometric algebra operators.
-Includes vector specifications, einsum pattern generation, and solvers.
+Includes vector specifications, einsum pattern generation, solvers, and contraction.
 """
 
+from morphis.algebra.contraction import (
+    IndexedTensor as IndexedTensor,
+    contract as contract,
+)
 from morphis.algebra.patterns import (
     INPUT_COLLECTION as INPUT_COLLECTION,
     INPUT_GEOMETRIC as INPUT_GEOMETRIC,

morphis/algebra/contraction.py

@@ -0,0 +1,263 @@
+"""
+Linear Algebra - Tensor Contraction
+
+Provides two contraction APIs for Morphis tensors:
+1. Index notation: G["mnab"] * q["n"] - bracket syntax with IndexedTensor
+2. Einsum-style: contract("mnab, n -> mab", G, q) - explicit signature
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from numpy import einsum
+
+
+if TYPE_CHECKING:
+    from morphis.elements.vector import Vector
+    from morphis.operations.operator import Operator
+
+
+# =============================================================================
+# IndexedTensor - Bracket Syntax API
+# =============================================================================
+
+
+class IndexedTensor:
+    """
+    Lightweight wrapper that pairs a tensor with index labels for contraction.
+
+    This class enables einsum-style syntax:
+        G["mnab"] * q["n"]  # contracts on index 'n'
+
+    The wrapper holds a reference (not a copy) to the underlying tensor,
+    making indexing O(1). Computation only happens when two IndexedTensors
+    are multiplied.
+
+    Attributes:
+        tensor: The underlying Vector or Operator (reference, not copy)
+        indices: String of index labels (e.g., "mnab")
+
+    Examples:
+        >>> G = Operator(...)  # lot=(M, N), grade=2 output
+        >>> q = Vector(...)  # lot=(N,), grade=0
+        >>> b = G["mnab"] * q["n"]  # contracts on 'n', result has indices "mab"
+    """
+
+    __slots__ = ("tensor", "indices")
+
+    def __init__(self, tensor: "Vector | Operator", indices: str):
+        """
+        Create an indexed tensor wrapper.
+
+        Args:
+            tensor: The underlying Vector or Operator
+            indices: String of index labels, one per axis of tensor.data
+        """
+        self.tensor = tensor
+        self.indices = indices
+
+        # Validate index count matches tensor dimensions
+        expected_ndim = tensor.data.ndim
+        if len(indices) != expected_ndim:
+            raise ValueError(
+                f"Index string '{indices}' has {len(indices)} indices, but tensor has {expected_ndim} dimensions"
+            )
+
+    def __mul__(self, other: "IndexedTensor") -> "Vector":
+        """
+        Contract two indexed tensors on matching indices.
+
+        Args:
+            other: Another IndexedTensor to contract with
+
+        Returns:
+            Vector with the contracted result
+        """
+        if not isinstance(other, IndexedTensor):
+            return NotImplemented
+
+        return _contract_indexed(self, other)
+
+    def __rmul__(self, other: "IndexedTensor") -> "Vector":
+        """Right multiplication for contraction."""
+        if not isinstance(other, IndexedTensor):
+            return NotImplemented
+
+        return _contract_indexed(other, self)
+
+    def __repr__(self) -> str:
+        tensor_type = type(self.tensor).__name__
+        return f"IndexedTensor({tensor_type}, indices='{self.indices}')"
+
+
+def _contract_indexed(*indexed_tensors: IndexedTensor) -> "Vector":
+    """
+    Contract multiple IndexedTensor objects.
+
+    Internal function that performs the actual contraction for bracket syntax.
+    """
+    from morphis.elements.vector import Vector
+
+    if len(indexed_tensors) < 2:
+        raise ValueError("Contraction requires at least 2 indexed tensors")
+
+    # Collect all index information
+    all_indices = [it.indices for it in indexed_tensors]
+    all_data = [it.tensor.data for it in indexed_tensors]
+
+    # Count index occurrences to determine output indices
+    index_counts: dict[str, int] = {}
+    for indices in all_indices:
+        for idx in indices:
+            index_counts[idx] = index_counts.get(idx, 0) + 1
+
+    # Output indices are those that appear exactly once (not contracted)
+    # Preserve order of first appearance
+    seen = set()
+    output_indices = ""
+    for indices in all_indices:
+        for idx in indices:
+            if idx not in seen:
+                seen.add(idx)
+                if index_counts[idx] == 1:
+                    output_indices += idx
+
+    # Build einsum signature
+    input_sig = ",".join(all_indices)
+    einsum_sig = f"{input_sig}->{output_indices}"
+
+    # Perform contraction
+    result_data = einsum(einsum_sig, *all_data)
+
+    # Get metric from first tensor that has one
+    metric = None
+    for it in indexed_tensors:
+        if hasattr(it.tensor, "metric") and it.tensor.metric is not None:
+            metric = it.tensor.metric
+            break
+
+    # Infer grade from output
+    result_grade = _infer_grade_from_indexed(indexed_tensors, output_indices)
+
+    return Vector(data=result_data, grade=result_grade, metric=metric)
+
+
+def _infer_grade_from_indexed(indexed_tensors: tuple[IndexedTensor, ...], output_indices: str) -> int:
+    """Infer grade for IndexedTensor contraction result."""
+    from morphis.elements.vector import Vector
+
+    # Track which indices are geometric (vs lot)
+    geo_indices = set()
+
+    for it in indexed_tensors:
+        if isinstance(it.tensor, Vector):
+            n_lot = len(it.tensor.lot)
+            n_geo = it.tensor.grade
+            # Geometric indices are the last 'grade' indices
+            geo_part = it.indices[n_lot : n_lot + n_geo]
+            geo_indices.update(geo_part)
+
+    # Count geometric indices in output
+    result_grade = sum(1 for idx in output_indices if idx in geo_indices)
+    return result_grade
+
+
+# =============================================================================
+# contract() - Einsum-Style API
+# =============================================================================
+
+
+def contract(signature: str, *tensors: "Vector | Operator") -> "Vector":
+    """
+    Einsum-style contraction for Morphis tensors.
+
+    Works exactly like numpy.einsum, but accepts Vector and Operator objects.
+    Extracts the underlying data, performs the einsum, and wraps the result
+    back into a Vector.
+
+    Args:
+        signature: Einsum signature string (e.g., "mn, n -> m")
+        *tensors: Morphis objects (Vector or Operator) to contract
+
+    Returns:
+        Vector containing the contracted result
+
+    Examples:
+        >>> g = euclidean_metric(3)
+        >>> u = Vector([1, 2, 3], grade=1, metric=g)
+        >>> v = Vector([4, 5, 6], grade=1, metric=g)
+
+        >>> # Dot product
+        >>> s = contract("a, a ->", u, v)
+        >>> s.data  # 1*4 + 2*5 + 3*6 = 32
+
+        >>> # Matrix-vector product
+        >>> M = Vector(data, grade=2, metric=g)  # shape (3, 3)
+        >>> w = contract("ab, b -> a", M, v)
+
+        >>> # Outer product
+        >>> outer = contract("a, b -> ab", u, v)
+
+        >>> # Batch contraction
+        >>> G = Vector(data, grade=2, lot=(M, N), metric=g)  # shape (M, N, 3, 3)
+        >>> q = Vector(data, grade=0, lot=(N,), metric=g)  # shape (N,)
+        >>> b = contract("mnab, n -> mab", G, q)
+    """
+    from morphis.elements.vector import Vector
+
+    if len(tensors) < 1:
+        raise ValueError("contract() requires at least 1 tensor")
+
+    # Extract data arrays from tensors
+    data_arrays = [t.data for t in tensors]
+
+    # Normalize signature: allow spaces around comma and arrow
+    sig = signature.replace(" ", "")
+
+    # Perform einsum
+    result_data = einsum(sig, *data_arrays)
+
+    # Get metric from first tensor that has one
+    metric = None
+    for t in tensors:
+        if hasattr(t, "metric") and t.metric is not None:
+            metric = t.metric
+            break
+
+    # Infer grade from output shape
+    result_grade = _infer_grade_from_signature(sig, tensors, result_data)
+
+    return Vector(data=result_data, grade=result_grade, metric=metric)
+
+
+def _infer_grade_from_signature(signature: str, tensors: tuple, result_data) -> int:
+    """Infer grade for einsum-style contraction result."""
+    from morphis.elements.vector import Vector
+
+    # Parse signature to get output indices
+    if "->" in signature:
+        input_part, output_indices = signature.split("->")
+    else:
+        # No explicit output - numpy determines it
+        return 0 if result_data.ndim == 0 else result_data.ndim
+
+    input_parts = input_part.split(",")
+
+    # Track which indices are geometric (vs lot)
+    geo_indices = set()
+
+    for k, t in enumerate(tensors):
+        if k < len(input_parts) and isinstance(t, Vector):
+            indices = input_parts[k]
+            n_lot = len(t.lot)
+            n_geo = t.grade
+            # Geometric indices are the last 'grade' indices
+            if len(indices) >= n_lot + n_geo:
+                geo_part = indices[n_lot : n_lot + n_geo]
+                geo_indices.update(geo_part)
+
+    # Count geometric indices in output
+    result_grade = sum(1 for idx in output_indices if idx in geo_indices)
+
+    return result_grade

morphis/algebra/patterns.py

@@ -5,13 +5,14 @@ Generates einsum signatures for linear operator operations. Uses disjoint index
 pools to avoid collisions between input and output indices.
 
 Index naming convention:
+- OUTPUT_LOT: "KLMN" (up to 4 output lot dims)
+- INPUT_LOT: "nopq" (up to 4 input lot dims)
 - OUTPUT_GEOMETRIC: "WXYZ" (up to grade-4 output blades)
-- OUTPUT_COLLECTION: "KLMN" (up to 4 output collection dims)
-- INPUT_COLLECTION: "nopq" (up to 4 input collection dims)
 - INPUT_GEOMETRIC: "abcd" (up to grade-4 input blades)
 
-Operator storage convention: (*output_geometric, *output_collection, *input_collection, *input_geometric)
-Vector storage convention: (*collection, *geometric)
+Storage conventions (lot-first, matching Vector layout):
+- Operator: (*out_lot, *in_lot, *out_geo, *in_geo)
+- Vector: (*lot, *geo)
 """
 
 from functools import lru_cache
@@ -20,67 +21,71 @@ from morphis.algebra.specs import VectorSpec
 
 
 # Index pools (disjoint to avoid collisions)
+OUTPUT_LOT = "KLMN"
+INPUT_LOT = "nopq"
 OUTPUT_GEOMETRIC = "WXYZ"
-OUTPUT_COLLECTION = "KLMN"
-INPUT_COLLECTION = "nopq"
 INPUT_GEOMETRIC = "abcd"
 
+# Backwards compatibility aliases
+OUTPUT_COLLECTION = OUTPUT_LOT
+INPUT_COLLECTION = INPUT_LOT
+
 
 @lru_cache(maxsize=128)
 def forward_signature(input_spec: VectorSpec, output_spec: VectorSpec) -> str:
     """
     Generate einsum signature for forward operator application: y = L * x
 
-    Operator data has shape: (*output_geometric, *output_collection, *input_collection, *input_geometric)
-    Input blade has shape: (*input_collection, *input_geometric)
-    Output blade has shape: (*output_collection, *output_geometric)
+    Operator data has shape: (*out_lot, *in_lot, *out_geo, *in_geo)
+    Input Vector has shape: (*in_lot, *in_geo)
+    Output Vector has shape: (*out_lot, *out_geo)
 
     Args:
-        input_spec: Specification of input blade
-        output_spec: Specification of output blade
+        input_spec: Specification of input Vector
+        output_spec: Specification of output Vector
 
     Returns:
-        Einsum signature string, e.g., "WXKn,n->KWX" for scalar->bivector
+        Einsum signature string, e.g., "KnWX,n->KWX" for scalar->bivector
 
     Examples:
         >>> # Scalar currents (N,) to bivector fields (M, 3, 3)
         >>> sig = forward_signature(
-        ...     VectorSpec(grade=0, collection=1, dim=3),
-        ...     VectorSpec(grade=2, collection=1, dim=3),
+        ...     VectorSpec(grade=0, lot=(1,), dim=3),
+        ...     VectorSpec(grade=2, lot=(1,), dim=3),
         ... )
         >>> sig
-        'WXKn,n->KWX'
+        'KnWX,n->KWX'
 
         >>> # Vector (N, 3) to bivector (M, 3, 3)
         >>> sig = forward_signature(
-        ...     VectorSpec(grade=1, collection=1, dim=3),
-        ...     VectorSpec(grade=2, collection=1, dim=3),
+        ...     VectorSpec(grade=1, lot=(1,), dim=3),
+        ...     VectorSpec(grade=2, lot=(1,), dim=3),
         ... )
         >>> sig
-        'WXKna,na->KWX'
+        'KnWXa,na->KWX'
     """
     _validate_spec_limits(input_spec, output_spec)
 
-    # Output geometric indices (stored first in operator)
-    out_geo = OUTPUT_GEOMETRIC[: output_spec.grade]
+    # Output lot indices
+    out_lot = OUTPUT_LOT[: output_spec.collection]
 
-    # Output collection indices
-    out_coll = OUTPUT_COLLECTION[: output_spec.collection]
+    # Input lot indices (contracted)
+    in_lot = INPUT_LOT[: input_spec.collection]
 
-    # Input collection indices (contracted)
-    in_coll = INPUT_COLLECTION[: input_spec.collection]
+    # Output geometric indices
+    out_geo = OUTPUT_GEOMETRIC[: output_spec.grade]
 
     # Input geometric indices (contracted)
     in_geo = INPUT_GEOMETRIC[: input_spec.grade]
 
-    # Build operator signature: out_geo + out_coll + in_coll + in_geo
-    op_indices = out_geo + out_coll + in_coll + in_geo
+    # Build operator signature: out_lot + in_lot + out_geo + in_geo (lot-first)
+    op_indices = out_lot + in_lot + out_geo + in_geo
 
-    # Build input signature: in_coll + in_geo (blade storage order)
-    input_indices = in_coll + in_geo
+    # Build input signature: in_lot + in_geo (Vector storage order)
+    input_indices = in_lot + in_geo
 
-    # Build output signature: out_coll + out_geo (blade storage order)
-    output_indices = out_coll + out_geo
+    # Build output signature: out_lot + out_geo (Vector storage order)
+    output_indices = out_lot + out_geo
 
     return f"{op_indices},{input_indices}->{output_indices}"
 
@@ -102,28 +107,28 @@ def adjoint_signature(input_spec: VectorSpec, output_spec: VectorSpec) -> str:
     Examples:
         >>> # Adjoint of scalar->bivector: bivector->scalar
         >>> sig = adjoint_signature(
-        ...     VectorSpec(grade=0, collection=1, dim=3),
-        ...     VectorSpec(grade=2, collection=1, dim=3),
+        ...     VectorSpec(grade=0, lot=(1,), dim=3),
+        ...     VectorSpec(grade=2, lot=(1,), dim=3),
         ... )
         >>> sig
-        'WXKn,KWX->n'
+        'KnWX,KWX->n'
     """
     _validate_spec_limits(input_spec, output_spec)
 
     # Same index allocation as forward
+    out_lot = OUTPUT_LOT[: output_spec.collection]
+    in_lot = INPUT_LOT[: input_spec.collection]
     out_geo = OUTPUT_GEOMETRIC[: output_spec.grade]
-    out_coll = OUTPUT_COLLECTION[: output_spec.collection]
-    in_coll = INPUT_COLLECTION[: input_spec.collection]
     in_geo = INPUT_GEOMETRIC[: input_spec.grade]
 
-    # Operator indices unchanged: out_geo + out_coll + in_coll + in_geo
-    op_indices = out_geo + out_coll + in_coll + in_geo
+    # Operator indices: out_lot + in_lot + out_geo + in_geo (lot-first)
+    op_indices = out_lot + in_lot + out_geo + in_geo
 
-    # Adjoint input (original output vec): out_coll + out_geo
-    adjoint_input = out_coll + out_geo
+    # Adjoint input (original output vec): out_lot + out_geo
+    adjoint_input = out_lot + out_geo
 
-    # Adjoint output (original input space): in_coll + in_geo
-    adjoint_output = in_coll + in_geo
+    # Adjoint output (original input space): in_lot + in_geo
+    adjoint_output = in_lot + in_geo
 
     return f"{op_indices},{adjoint_input}->{adjoint_output}"
 
@@ -131,42 +136,52 @@ def adjoint_signature(input_spec: VectorSpec, output_spec: VectorSpec) -> str:
 def operator_shape(
     input_spec: VectorSpec,
     output_spec: VectorSpec,
-    input_collection: tuple[int, ...],
-    output_collection: tuple[int, ...],
+    input_lot: tuple[int, ...] | None = None,
+    output_lot: tuple[int, ...] | None = None,
+    *,
+    input_collection: tuple[int, ...] | None = None,
+    output_collection: tuple[int, ...] | None = None,
 ) -> tuple[int, ...]:
     """
-    Compute the expected shape of operator data given specs and collection shapes.
+    Compute the expected shape of operator data given specs and lot shapes.
 
     Args:
-        input_spec: Specification of input blade
-        output_spec: Specification of output blade
-        input_collection: Shape of input collection dimensions
-        output_collection: Shape of output collection dimensions
+        input_spec: Specification of input Vector
+        output_spec: Specification of output Vector
+        input_lot: Shape of input lot dimensions
+        output_lot: Shape of output lot dimensions
+        input_collection: DEPRECATED alias for input_lot
+        output_collection: DEPRECATED alias for output_lot
 
     Returns:
-        Operator data shape: (*output_geometric, *output_collection, *input_collection, *input_geometric)
+        Operator data shape: (*out_lot, *in_lot, *out_geo, *in_geo)
 
     Examples:
         >>> # Scalar (N=5) to bivector (M=10) in 3D
         >>> shape = operator_shape(
-        ...     VectorSpec(grade=0, collection=1, dim=3),
-        ...     VectorSpec(grade=2, collection=1, dim=3),
-        ...     input_collection=(5,),
-        ...     output_collection=(10,),
+        ...     VectorSpec(grade=0, lot=(1,), dim=3),
+        ...     VectorSpec(grade=2, lot=(1,), dim=3),
+        ...     input_lot=(5,),
+        ...     output_lot=(10,),
         ... )
         >>> shape
-        (3, 3, 10, 5)
+        (10, 5, 3, 3)
     """
-    if len(input_collection) != input_spec.collection:
-        raise ValueError(
-            f"input_collection has {len(input_collection)} dims, but input_spec expects {input_spec.collection}"
-        )
-    if len(output_collection) != output_spec.collection:
-        raise ValueError(
-            f"output_collection has {len(output_collection)} dims, but output_spec expects {output_spec.collection}"
-        )
+    # Handle deprecated aliases
+    if input_lot is None and input_collection is not None:
+        input_lot = input_collection
+    if output_lot is None and output_collection is not None:
+        output_lot = output_collection
+
+    if input_lot is None or output_lot is None:
+        raise ValueError("input_lot and output_lot are required")
+
+    if len(input_lot) != input_spec.collection:
+        raise ValueError(f"input_lot has {len(input_lot)} dims, but input_spec expects {input_spec.collection}")
+    if len(output_lot) != output_spec.collection:
+        raise ValueError(f"output_lot has {len(output_lot)} dims, but output_spec expects {output_spec.collection}")
 
-    return output_spec.geometric_shape + output_collection + input_collection + input_spec.geometric_shape
+    return output_lot + input_lot + output_spec.geo + input_spec.geo
 
 
 def _validate_spec_limits(input_spec: VectorSpec, output_spec: VectorSpec) -> None:
@@ -175,7 +190,7 @@ def _validate_spec_limits(input_spec: VectorSpec, output_spec: VectorSpec) -> No
         raise ValueError(f"Input grade {input_spec.grade} exceeds index pool limit {len(INPUT_GEOMETRIC)}")
     if output_spec.grade > len(OUTPUT_GEOMETRIC):
         raise ValueError(f"Output grade {output_spec.grade} exceeds index pool limit {len(OUTPUT_GEOMETRIC)}")
-    if input_spec.collection > len(INPUT_COLLECTION):
-        raise ValueError(f"Input collection {input_spec.collection} exceeds limit {len(INPUT_COLLECTION)}")
-    if output_spec.collection > len(OUTPUT_COLLECTION):
-        raise ValueError(f"Output collection {output_spec.collection} exceeds limit {len(OUTPUT_COLLECTION)}")
+    if input_spec.collection > len(INPUT_LOT):
+        raise ValueError(f"Input lot dims {input_spec.collection} exceeds limit {len(INPUT_LOT)}")
+    if output_spec.collection > len(OUTPUT_LOT):
+        raise ValueError(f"Output lot dims {output_spec.collection} exceeds limit {len(OUTPUT_LOT)}")

morphis/algebra/solvers.py

@@ -30,7 +30,7 @@ def _to_matrix(op: "Operator") -> NDArray:
     """
     Flatten operator to matrix form for linear algebra operations.
 
-    Reshapes from (*out_geo, *out_coll, *in_coll, *in_geo) to (out_flat, in_flat).
+    Reshapes from (*out_lot, *in_lot, *out_geo, *in_geo) to (out_flat, in_flat).
 
     Args:
         op: Operator to flatten
@@ -41,22 +41,20 @@ def _to_matrix(op: "Operator") -> NDArray:
     out_flat = int(prod(op.output_shape))
     in_flat = int(prod(op.input_shape))
 
-    # We need to reorder axes to group output together and input together
-    # Current: (*out_geo, *out_coll, *in_coll, *in_geo)
-    # Target for reshape: (*out_coll, *out_geo, *in_coll, *in_geo)
-    # But we want blade order: (*coll, *geo), so:
-    # Target: (*out_coll, *out_geo, *in_coll, *in_geo)
-
-    out_geo_axes = list(range(op.output_spec.grade))
-    out_coll_start = op.output_spec.grade
-    out_coll_axes = list(range(out_coll_start, out_coll_start + op.output_spec.collection))
-    in_coll_start = out_coll_start + op.output_spec.collection
-    in_coll_axes = list(range(in_coll_start, in_coll_start + op.input_spec.collection))
-    in_geo_start = in_coll_start + op.input_spec.collection
+    # Current lot-first layout: (*out_lot, *in_lot, *out_geo, *in_geo)
+    # Target for reshape: (*out_lot, *out_geo, *in_lot, *in_geo)
+    # This groups output dims together (lot, geo) and input dims together (lot, geo)
+
+    out_lot_axes = list(range(op.output_spec.collection))
+    in_lot_start = op.output_spec.collection
+    in_lot_axes = list(range(in_lot_start, in_lot_start + op.input_spec.collection))
+    out_geo_start = in_lot_start + op.input_spec.collection
+    out_geo_axes = list(range(out_geo_start, out_geo_start + op.output_spec.grade))
+    in_geo_start = out_geo_start + op.output_spec.grade
     in_geo_axes = list(range(in_geo_start, in_geo_start + op.input_spec.grade))
 
-    # Reorder to: (*out_coll, *out_geo, *in_coll, *in_geo)
-    perm = out_coll_axes + out_geo_axes + in_coll_axes + in_geo_axes
+    # Reorder to: (*out_lot, *out_geo, *in_lot, *in_geo)
+    perm = out_lot_axes + out_geo_axes + in_lot_axes + in_geo_axes
     reordered = op.data.transpose(perm)
 
     return reordered.reshape(out_flat, in_flat)
@@ -86,22 +84,22 @@ def _from_matrix(
     """
     from morphis.operations.operator import Operator
 
-    # Reshape from (out_flat, in_flat) to (*out_coll, *out_geo, *in_coll, *in_geo)
+    # Reshape from (out_flat, in_flat) to (*out_lot, *out_geo, *in_lot, *in_geo)
     intermediate_shape = output_collection + output_spec.geometric_shape + input_collection + input_spec.geometric_shape
     tensor = matrix.reshape(intermediate_shape)
 
-    # Reorder from (*out_coll, *out_geo, *in_coll, *in_geo)
-    # to (*out_geo, *out_coll, *in_coll, *in_geo)
-    out_coll_axes = list(range(output_spec.collection))
+    # Reorder from (*out_lot, *out_geo, *in_lot, *in_geo)
+    # to (*out_lot, *in_lot, *out_geo, *in_geo)
+    out_lot_axes = list(range(output_spec.collection))
     out_geo_start = output_spec.collection
     out_geo_axes = list(range(out_geo_start, out_geo_start + output_spec.grade))
-    in_coll_start = out_geo_start + output_spec.grade
-    in_coll_axes = list(range(in_coll_start, in_coll_start + input_spec.collection))
-    in_geo_start = in_coll_start + input_spec.collection
+    in_lot_start = out_geo_start + output_spec.grade
+    in_lot_axes = list(range(in_lot_start, in_lot_start + input_spec.collection))
+    in_geo_start = in_lot_start + input_spec.collection
     in_geo_axes = list(range(in_geo_start, in_geo_start + input_spec.grade))
 
-    # Target order: (*out_geo, *out_coll, *in_coll, *in_geo)
-    perm = out_geo_axes + out_coll_axes + in_coll_axes + in_geo_axes
+    # Target order: (*out_lot, *in_lot, *out_geo, *in_geo)
+    perm = out_lot_axes + in_lot_axes + out_geo_axes + in_geo_axes
     data = tensor.transpose(perm)
 
     return Operator(
@@ -266,24 +264,25 @@ def structured_svd(
     dim = op.output_spec.dim
 
     # U maps from reduced space (r,) to output space
-    u_input_spec = VectorSpec(grade=0, collection=1, dim=dim)
+    u_input_spec = VectorSpec(grade=0, lot=(1,), dim=dim)
 
     # Vt maps from input space to reduced space (r,)
-    vt_output_spec = VectorSpec(grade=0, collection=1, dim=dim)
+    vt_output_spec = VectorSpec(grade=0, lot=(1,), dim=dim)
 
     # Wrap U as Operator
     # U_mat has shape (out_flat, r)
-    # Need to reshape to (*out_coll, *out_geo, r)
-    # Then reorder to (*out_geo, *out_coll, r)
+    # Reshape to (*out_lot, *out_geo, r) then reorder to lot-first: (*out_lot, r, *out_geo)
+    # out_flat = prod(out_lot) * prod(out_geo), input is (r,) which is just lot dims
     U_intermediate = U_mat.reshape(op.output_collection + op.output_spec.geometric_shape + (r,))
 
-    out_coll_axes = list(range(op.output_spec.collection))
+    out_lot_axes = list(range(op.output_spec.collection))
     out_geo_start = op.output_spec.collection
     out_geo_axes = list(range(out_geo_start, out_geo_start + op.output_spec.grade))
     r_axis = [out_geo_start + op.output_spec.grade]
 
-    # Target: (*out_geo, *out_coll, r)
-    U_perm = out_geo_axes + out_coll_axes + r_axis
+    # Target lot-first: (*out_lot, r, *out_geo, no_in_geo)
+    # Since in_spec is grade=0, no in_geo axes
+    U_perm = out_lot_axes + r_axis + out_geo_axes
     U_data = U_intermediate.transpose(U_perm)
 
     U_op = Operator(
@@ -295,13 +294,13 @@ def structured_svd(
 
     # Wrap Vt as Operator
     # Vt_mat has shape (r, in_flat)
-    # Need to reshape to (r, *in_coll, *in_geo)
+    # Reshape to (r, *in_lot, *in_geo)
     Vt_intermediate = Vt_mat.reshape((r,) + op.input_collection + op.input_spec.geometric_shape)
 
-    # Current order: (r, *in_coll, *in_geo)
-    # Target order for operator data: (*out_geo, *out_coll, *in_coll, *in_geo)
-    # For Vt, out is reduced (grade=0, coll_dims=1) so no out_geo, out_coll = (r,)
-    # So target: (r, *in_coll, *in_geo) which is already correct!
+    # Current order: (r, *in_lot, *in_geo)
+    # Target lot-first: (*out_lot, *in_lot, *out_geo, *in_geo)
+    # For Vt, out is (r,) so out_lot=(r,) and out_geo=() since grade=0
+    # So target: (r, *in_lot, *in_geo) which is already correct!
     Vt_data = Vt_intermediate
 
     Vt_op = Operator(