corewise_stack_dot#

t3toolbox.corewise.corewise_stack_dot(X, Y, n_stack)#
def corewise_stack_dot(
        X:          NDArrayTree,  # tree; each leaf shape = stack_shape + (core dims)
        Y:          NDArrayTree,  # same tree structure / leaf shapes as X
        n_stack:    int,          # number of leading stack axes kept (contract the rest)
) -> NDArray:  # array of shape = common stack_shape (scalar when n_stack==0)

Like corewise_dot, but vectorized over the leading n_stack (stack) axes.

Each leaf is contracted over its trailing (non-stack) axes only, keeping the leading n_stack axes, and the per-leaf results are summed. Returns an array of shape equal to the common leading stack shape (a scalar when n_stack == 0, matching corewise_dot()).

Examples

>>> import numpy as np
>>> import t3toolbox.corewise as cw
>>> X = (np.ones((2, 3)), np.ones((2, 4, 5)))   # stack=(2,), then core axes
>>> Y = (2*np.ones((2, 3)), np.ones((2, 4, 5)))
>>> print(cw.corewise_stack_dot(X, Y, 1))       # per-stack-slice: 2*3 + 4*5 = 26
[26. 26.]
Parameters:
Return type:

NDArray