corewise_stack_sum#

t3toolbox.corewise.corewise_stack_sum(X, axis, n_stack)#
def corewise_stack_sum(
        X:          NDArrayTree,  # tree; each leaf shape = stack_shape + (core dims)
        axis,                     # stack axis/axes (None -> all; negatives wrap rel. to n_stack)
        n_stack:    int,          # number of leading stack axes axis is normalized against
) -> NDArrayTree:                 # same tree structure as X, summed stack axes removed

Sum each leaf over stack axes, vectorized.

Normalizes axis against the n_stack leading (stack) axes (None -> all of them; negative axes wrap relative to n_stack), then corewise_sum() over those axes.

Examples

>>> import numpy as np
>>> import t3toolbox.corewise as cw
>>> X = (np.ones((2, 3, 4)),)   # stack=(2, 3), then a core axis
>>> print(cw.corewise_stack_sum(X, None, 2)[0].shape)   # sum both stack axes
(4,)
Parameters:
Return type:

NDArrayTree