Added MomentumBlockTensor and compatibilities upgrades

[Alphaharrius]

Summary

This PR extends the momentum-block workflow so Fourier and band-space helpers can work cleanly with asymmetric left/right Hilbert spaces and block-space momentum projections.

It also introduces a stricter subtype policy for structured tensor wrappers like MomentumBlockTensor: layout-preserving operations keep the subtype, while rank/layout-changing generic tensor operations downgrade to plain Tensor.

What changed

One-sided momentum-block band helpers

Updated the band-transform and band-fold factorization APIs to work one side at a time:

• get_band_transform(t, tensor, side="left")
• get_band_fold(transform, tensor, side="left")
• bandtransform(t, tensor, opt="left" | "right" | "both")
• bandfold(transform, tensor, opt="left" | "right" | "both")

These helpers no longer require the last two Hilbert-space legs to span the same rays. Instead, the selected side determines which Hilbert-space leg is sampled when constructing the block transform/fold.

This makes rectangular and left/right-asymmetric band tensors first-class in the factorized workflow.

region_restrict(..., side=...) and block-space momentum projection

Extended region_restrict to accept:

side: Literal["left", "right"] = "left"

for both runtime and stub signatures.

Behavior now is:

• for ordinary rank-3 Fourier/band tensors, side="left" uses tensor.dims[1] and side="right" uses tensor.dims[2]
• if the leading axis is a MomentumBlockSpace, region_restrict projects that pair axis to a unique MomentumSpace through:
  • MomentumBlockSpace.left()
  • MomentumBlockSpace.right()

The return value is still a plain Fourier Tensor, which is exactly the object needed for products such as:

F @ MBT
MBT @ F.h(-2, -1)

To support that, MomentumBlockSpace now provides:

• left() -> MomentumSpace
• right() -> MomentumSpace

and also defines a proper structure-based hash so these spaces work reliably in cached alignment/union code.

Strict-dims tensor subtype policy

Added a new decorator and runtime policy in qten.linalg.tensors:

@strict_dims

Subclasses marked with @strict_dims declare that their symbolic dimension layout is part of their semantic invariant.

For these subclasses:

• layout-preserving generic operations may keep the subtype
• rank/layout-changing generic operations automatically downgrade to plain Tensor

This is implemented through a shared helper:

• _wrap_tensor_result(...)

and a class marker:

• Tensor._strict_dims: ClassVar[bool] = False

MomentumBlockTensor is now marked with @strict_dims.

This fixes the previous issue where generic methods like mean(dim=0) attempted to preserve MomentumBlockTensor even after destroying its defining layout:

(MomentumBlockSpace, HilbertSpace, HilbertSpace)

The generic downgrade behavior was applied to rank/layout-changing operations such as:

• unsqueeze
• squeeze
• all
• mean
• norm
• argmax
• argmin
• factorize_dim
• product_dims

Tensor addition preserves left subtype

Updated generic Tensor + Tensor so it preserves the concrete type of the left operand by using replace(left, ...) instead of always constructing a fresh base Tensor.

This keeps subtype-preserving additions working naturally for structured wrappers whose output dims still satisfy their invariants.

The tensors.pyi stubs were updated accordingly for __add__.

Documentation and stubs

Expanded and updated documentation for:

• MomentumBlockTensor
• get_band_transform
• get_band_fold
• bandtransform
• bandfold
• region_restrict
• MomentumBlockSpace.left/right
• strict_dims
• Tensor._strict_dims

Also updated the relevant .pyi files for:

• fourier.region_restrict
• MomentumBlockSpace.left/right
• Tensor.__add__
• strict_dims
• Tensor._strict_dims

Tests

Updated and added tests covering:

• one-sided factorization equivalence for bandtransform
• one-sided factorization equivalence for bandfold
• left/right sampled-side behavior for block transforms/folds
• support for asymmetric left/right Hilbert spaces
• MomentumBlockSpace projection behavior through region_restrict
• MomentumBlockTensor subtype preservation through generic addition
• strict-dims downgrade behavior for:
  • mean
  • unsqueeze
  • argmax

Notes

Two important semantic points in this PR:

• region_restrict(MomentumBlockTensor, ...) still returns a plain Fourier Tensor, not a block tensor. That is intentional: it is meant to produce the left/right Fourier factor that can act on a block tensor.

Validation

Validated with:

• uv run pytest -q tests/test_momentum_block_tensor.py tests/test_tensors.py
• compile checks for the touched modules via python -m py_compile

[Copilot]

Pull request overview

This PR introduces momentum-pair “block” tensor infrastructure to factorize band-space transforms/folds via explicit momentum-block operators, and updates the band helper APIs and tests accordingly.

Changes:

• Added MomentumBlockTensor (rank-3 structured tensor) with specialized permute/transpose and @ rules for momentum-pair block contraction/composition.
• Added MomentumBlockSpace (momentum-pair state space) and updated bands.bandtransform/bandfold to use get_band_transform/get_band_fold + block-conjugation factoring.
• Added/updated tests to cover factorization equivalence and momentum-block matmul semantics; bumped package version to 0.4.0.

Reviewed changes

Copilot reviewed 14 out of 14 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
src/qten/linalg/_mb_tensor.py	New MomentumBlockTensor implementation + specialized @/axis-reordering semantics.
src/qten/bands.py	Adds get_band_transform/get_band_fold and refactors bandtransform/bandfold to use them.
src/qten/symbolics/state_space.py	Adds MomentumBlockSpace runtime implementation.
src/qten/symbolics/state_space.pyi	Adds MomentumBlockSpace stub.
src/qten/linalg/tensors.py	Makes permute/transpose multimethods to allow subtype-specific behavior.
src/qten/symbolics/__init__.py	Re-exports MomentumBlockSpace.
src/qten/__init__.py	Re-exports MomentumBlockTensor.
tests/test_momentum_block_tensor.py	New tests for transpose and momentum-block @ behavior.
tests/test_basis_transform.py	Adds tests for get_band_fold factorization and side sampling.
tests/test_abelian.py	Adds tests for get_band_transform factorization and side sampling.
pyproject.toml	Version bump to 0.4.0.

Comments suppressed due to low confidence (1)

src/qten/linalg/tensors.py:3199

• permute()/transpose() were switched to @multimethod and now have runtime signatures returning plain Tensor, but their docstrings still claim they return TensorType (preserving wrapper type). This also conflicts with the Tensor.permute/transpose methods being typed as -> Self and the stubs in tensors.pyi still declaring TensorType. Consider updating the annotations/docstrings (or adding overloads/casts) so the public typing contract continues to reflect subtype-preserving behavior.

    right : Tensor
        The tensor.

    Returns
    -------
    Tensor
        A new tensor with each element multiplied by the scalar.
    """
    return Tensor(data=left * right.data, dims=right.dims)


@Operable.__mul__.register
def _(left: Tensor, right: Number) -> Tensor:
    """
    Perform element-wise multiplication of a tensor and a number.

    Parameters
    ----------
    left : Tensor
        The tensor.
    right : Number
        The scalar value.

    Returns
    -------
    Tensor
        A new tensor with each element multiplied by the scalar.
    """
    return Tensor(data=left.data * right, dims=left.dims)


@Operable.__add__.register
def _(left: Number, right: Tensor) -> Tensor:
    r"""
    Add a number to the diagonal of the tensor (broadcasting over batch dimensions).

    This treats the tensor as a batch of matrices defined by the last two
    dimensions. The scalar is added to the diagonal elements of each matrix.
    For rank-2 tensors this is equivalent to \(cI + M\).

    Parameters
    ----------
    left : Number
        The scalar value to add to the diagonal.
    right : Tensor
        The target tensor (must be at least rank 2).

    Returns
    -------
    Tensor
        The result of adding the scalar to the diagonal.
    """
    iden = eye(right.dims)
    return left * iden + right


@Operable.__add__.register
def _(left: Tensor, right: Number) -> Tensor:
    r"""
    Add a number to the diagonal of the tensor (broadcasting over batch dimensions).

    This treats the tensor as a batch of matrices defined by the last two
    dimensions. The scalar is added to the diagonal elements of each matrix.
    For rank-2 tensors this is equivalent to \(M + cI\).

    Parameters
    ----------
    left : Tensor
        The target tensor (must be at least rank 2).
    right : Number
        The scalar value to add to the diagonal.

    Returns

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 14 out of 14 changed files in this pull request and generated 2 comments.

[Copilot]

Pull request overview

Copilot reviewed 14 out of 14 changed files in this pull request and generated 6 comments.

[Copilot]

Pull request overview

Copilot reviewed 15 out of 15 changed files in this pull request and generated 4 comments.

[eyjafjallac]

Nice

[Copilot]

Pull request overview

Copilot reviewed 15 out of 15 changed files in this pull request and generated 5 comments.

[Copilot]

Pull request overview

Copilot reviewed 16 out of 16 changed files in this pull request and generated 2 comments.