contract
========

.. py:function:: t3toolbox.backend.contractions.contract(subscripts, *operands, **group_lens)

   .. code-block:: python

      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)                                # doctest: +IGNORE_EXCEPTION_DETAIL
   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)