contract#

t3toolbox.backend.contractions.contract(subscripts, *operands, **group_lens)#
def contract(
        subscripts: str,        # grouped einsum, e.g. 'WCa,Caib,WCi->WCb': UPPERCASE = a group of
                                # ZERO OR MORE axes; lowercase = one axis. Explicit '->' required.
        *operands:  NDArray,    # one array per comma-separated input term
        **group_lens: int,      # len_<G> (e.g. len_W=2) = axis count of group <G>. Required only
                                # when the subscripts cannot pin it; verified when redundant.
) -> NDArray:

Grouped einsum: an einsum whose UPPERCASE letters each stand for a GROUP of zero or more axes.

The group sizes are solved from the operand ndims, every group is expanded into ordinary single-axis letters, and one einsum runs on the operands exactly as given – no reshape, no data movement, every axis of every group independently shardable. Dispatch follows the house convention (inferred numpy/jax; numpy uses the greedy pairwise BLAS path, jax one fused einsum).

Full usage guide with examples: docs/grouped_contractions.md; block-letter conventions (W probe stack, K tangent stack, C frame stack, frame-inner order) are in docs/batching_and_stacking.md; implementation internals: docs/contributor/contractions_internals.md. Group axes follow einsum’s usual extent rules: mismatched sizes for the same group axis raise, except that a size-1 axis broadcasts (standard einsum semantics on both backends).

A batched matvec where the matrix batch C is shared and the vector batch W rides:

>>> import numpy as np
>>> from t3toolbox.backend.contractions import contract
>>> A = np.random.default_rng(0).standard_normal((2, 4, 3))         # C=(2,), axes (i, o)
>>> x = np.random.default_rng(1).standard_normal((5, 2, 3))         # W=(5,), C=(2,), axis o
>>> y = contract('Cio,WCo->WCi', A, x)
>>> y.shape
(5, 2, 4)
>>> bool(np.allclose(y, np.einsum('cio,wco->wci', A, x)))           # |W|=|C|=1 by hand
True

The SAME string handles any group sizes, including empty (a group of zero axes just vanishes):

>>> contract('Cio,WCo->WCi', np.ones((4, 3)), np.ones(3)).shape     # W=(), C=()
(4,)
>>> contract('Cio,WCo->WCi', np.ones((2, 2, 4, 3)), np.ones((5, 6, 2, 2, 3))).shape
(5, 6, 2, 2, 4)

When the ndims cannot pin a split, say which axes are which via len_<G> – the error says exactly what is missing (and this is decided from the subscripts alone, never from the shapes, so a call site either always needs the argument or never does):

>>> zo = np.ones((5, 2, 3)); za = np.ones((5, 2, 6))                # W=(5,), C=(2,)
>>> contract('WCo,WCa->Cao', zo, za)
Traceback (most recent call last):
    ...
ValueError: contract('WCo,WCa->Cao'): the operand ndims do not determine ...
>>> contract('WCo,WCa->Cao', zo, za, len_W=1).shape                 # sum over the probe stack W
(2, 6, 3)
Parameters:

  • subscripts (str)

  • operands (NDArray)

  • group_lens (int)

Return type:

NDArray