Следующая функция позволяет уменьшить заданную ось с различными срезами, указанными массивами начала и остановки. Он использует np.ufunc.reduceat под капотом вместе с соответствующим образом измененными версиями входного массива и индексов. Это позволяет избежать ненужных вычислений, но выделяет промежуточный массив, в два раза превышающий размер конечного выходного массива (однако вычисление отброшенных значений не выполняется).

def sliced_reduce(a, i, j, ufunc, axis=None):
    """Reduce an array along a given axis for varying slices `a[..., i:j, ...]` where `i` and `j` are arrays themselves.

    Parameters
    ----------
    a : array
        The array to be reduced.
    i : array
        Start indices for the reduced axis. Must have the same shape as `j`.
    j : array
        Stop indices for the reduced axis. Must have the same shape as `i`.
    ufunc : function
        The function used for reducing the indicated axis.
    axis : int, optional
        Axis to be reduced. Defaults to `len(i.shape)`.

    Returns
    -------
    array
        Shape `i.shape + a.shape[axis+1:]`.

    Notes
    -----
    The shapes of `a` and `i`, `j` must match up to the reduced axis.
    That means `a.shape[:axis] == i.shape[len(i.shape) - axis:]``. 
    `i` and `j` can have additional leading dimensions and `a` can have additional trailing dimensions.
    """
    if axis is None:
        axis = len(i.shape)
    indices = np.tile(
        np.repeat(
            np.arange(np.prod(a.shape[:axis])) * a.shape[axis],
            2  # Repeat two times to have start and stop indices next to each other.
        ),
        np.prod(i.shape[:len(i.shape) - axis])  # Perform summation for each element of additional axes.
    )
    # Add `a.shape[axis]` to account for negative indices.
    indices[::2] += (a.shape[axis] + i.ravel()) % a.shape[axis]
    indices[1::2] += (a.shape[axis] + j.ravel()) % a.shape[axis]
    # Now indices are sorted in ascending order but this will lead to unnecessary computation when reducing
    # from odd to even indices (since we're only interested in even to odd indices).
    # Hence we reverse the order of index pairs (need to reverse the result as well then).
    indices = indices.reshape(-1, 2)[::-1].ravel()
    result = ufunc.reduceat(a.reshape(-1, *a.shape[axis+1:]), indices)[::2]  # Select only even to odd.
    # In case start and stop index are equal (i.e. empty slice) `reduceat` will select the element
    # corresponding to the start index. Need to supply the correct default value in this case.
    result[indices[::2] == indices[1::2]] = ufunc.reduce([])
    return result[::-1].reshape(*(i.shape + a.shape[axis+1:]))  # Reverse order and reshape.

Для примеров в ОП его можно использовать следующим образом:

# 1. example:
b = np.random.randint(0, 1000, size=(3, 10, 3))
v = np.random.randint(-9, 10, size=3)  # Indexing into `b.shape[1]`.
result = sliced_reduce(b, np.zeros_like(v), v, np.add)

# 2. example:
b = np.random.randint(0, 1000, size=(98, 3, 300, 3))
v = np.random.randint(-299, 300, size=(32, 98, 3))  # Indexing into `b.shape[2]`; one additional leading dimension for `v`.
result = sliced_reduce(b, np.zeros_like(v), v, np.add, axis=2)

Примечания

• Изменение порядка пар плоских индексов, чтобы иметь even < odd и, таким образом, сократить каждое второе вычисление с помощью no-op, не кажется хорошей идеей (вероятно, потому, что сглаженный массив больше не проходится в порядке расположения памяти). Удаление этой части и использование плоских индексов в порядке возрастания дает прирост производительности примерно на 30% (также для перфплоты, хотя и не включенного туда).

Для чего это стоит, вот один лайнер. Никаких обещаний, что это самая эффективная версия, потому что она делает гораздо больше дополнений, чем необходимо:

In [25]: b.cumsum(axis=1)[np.arange(b.shape[0]), v-1]                                                          
Out[25]: 
array([[3., 1., 4.],
       [4., 2., 3.],
       [3., 0., 1.]])

(Также имейте в виду, что он неправильно обрабатывает 0 в v.)

Следующая функция позволяет суммировать заданную ось с различными срезами, указанными массивами начала и остановки. Он использует np.einsum под капотом вместе с соответствующим образом вычисленным массивом коэффициентов, который указывает, какие элементы входного массива должны участвовать в сумме (с использованием коэффициентов 1 и 0). Использование einsum делает реализацию совместимой с другими пакетами, такими как ПиТорч или ТензорФлоу (с небольшими изменениями). Это удваивает количество необходимых вычислений, так как на каждую операцию сложения приходится дополнительная операция умножения с массивом коэффициентов.

from string import ascii_lowercase as symbols
import numpy as np

def sliced_sum(a, i, j, axis=None):
    """Sum an array along a given axis for varying slices `a[..., i:j, ...]` where `i` and `j` are arrays themselves.

    Parameters
    ----------
    a : array
        The array to be summed over.
    i : array
        The start indices for the summation axis. Must have the same shape as `j`.
    j : array
        The stop indices for the summation axis. Must have the same shape as `i`.
    axis : int, optional
        Axis to be summed over. Defaults to `len(i.shape)`.

    Returns
    -------
    array
        Shape `i.shape + a.shape[axis+1:]`.

    Notes
    -----
    The shapes of `a` and `i`, `j` must match up to the summation axis.
    That means `a.shape[:axis] == i.shape[len(i.shape) - axis:]``. 
    `i` and `j` can have additional leading dimensions and `a` can have additional trailing dimensions.
    """
    if axis is None:
        axis = len(i.shape)

    # Compute number of leading, common and trailing dimensions.
    l = len(i.shape) - axis      # Number of leading dimensions.
    m = len(i.shape) - l         # Number of common dimensions.
    n = len(a.shape) - axis - 1  # Number of trailing dimensions.

    # Select the corresponding symbols for `np.einsum`.
    leading = symbols[:l]
    common = symbols[l:l+m]
    summation = symbols[l+m]
    trailing = symbols[l+m+1:l+m+1+n]

    # Convert negative indices.
    i = (a.shape[axis] + i) % a.shape[axis]
    j = (a.shape[axis] + j) % a.shape[axis]

    # Compute the "active" elements, i.e. the ones that should participate in the summation.
    # "active" elements have a coefficient of 1 (True), others are 0 (False).
    indices, i, j = np.broadcast_arrays(np.arange(a.shape[axis]),
                                        np.expand_dims(i, -1), np.expand_dims(j, -1))
    active_elements = (i <= indices) & (indices < j)
    return np.einsum(f'{leading + common + summation},{common + summation + trailing}->{leading + common + trailing}',
                     active_elements, a)

Для примеров в ОП его можно использовать следующим образом:

# 1. example:
b = np.random.randint(0, 1000, size=(3, 10, 3))
v = np.random.randint(-9, 10, size=3)  # Indexing into `b.shape[1]`.
result = sliced_sum(b, np.zeros_like(v), v)

# 2. example:
b = np.random.randint(0, 1000, size=(98, 3, 300, 3))
v = np.random.randint(-299, 300, size=(32, 98, 3))  # Indexing into `b.shape[2]`; one additional leading dimension for `v`.
result = sliced_sum(b, np.zeros_like(v), v, axis=2)

Другой вариант — использовать Нумба для ускорения цикла. Это позволяет избежать ненужных вычислений и выделения памяти и полностью совместимо со всеми numpy функциями (т. е. также prod и т. д. работают аналогично).

import numba
import numpy as np

def sliced_sum_numba(a, i, j, axis=None):
    """Sum an array along a given axis for varying slices `a[..., i:j, ...]` where `i` and `j` are arrays themselves.

    Parameters
    ----------
    a : array
        The array to be summed over.
    i : array
        The start indices for the summation axis. Must have the same shape as `j`.
    j : array
        The stop indices for the summation axis. Must have the same shape as `i`.
    axis : int, optional
        Axis to be summed over. Defaults to `len(i.shape)`.

    Returns
    -------
    array
        Shape `i.shape + a.shape[axis+1:]`.

    Notes
    -----
    The shapes of `a` and `i`, `j` must match up to the summation axis.
    That means `a.shape[:axis] == i.shape[len(i.shape) - axis:]``. 
    `i` and `j` can have additional leading dimensions and `a` can have additional trailing dimensions.
    """
    if axis is None:
        axis = len(i.shape)
    # Convert negative indices.
    i = (a.shape[axis] + i) % a.shape[axis]
    j = (a.shape[axis] + j) % a.shape[axis]
    # Operate on a flattened version of the array (dimensions up to `axis` are flattened).
    m = np.prod(i.shape[:len(i.shape) - axis], dtype=int)  # Elements in leading dimensions.
    n = np.prod(i.shape[len(i.shape) - axis:], dtype=int)  # Elements in common dimensions.
    a_flat = a.reshape(-1, *a.shape[axis:])
    i_flat = i.ravel()
    j_flat = j.ravel()
    result = np.empty((m*n,) + a.shape[axis+1:], dtype=a.dtype)
    numba_sum(a_flat, i_flat, j_flat, m, n, result)
    return result.reshape(*(i.shape + a.shape[axis+1:]))

@numba.jit(parallel=True, nopython=True)
def numba_sum(a, i, j, m, n, out):
    for index in numba.prange(m*n):
        out[index] = np.sum(a[index % n, i[index]:j[index]], axis=0)

Для примеров в ОП его можно использовать следующим образом:

# 1. example:
b = np.random.randint(0, 1000, size=(3, 10, 3))
v = np.random.randint(-9, 10, size=3)  # Indexing into `b.shape[1]`.
result = sliced_sum_numba(b, np.zeros_like(v), v)

# 2. example:
b = np.random.randint(0, 1000, size=(98, 3, 300, 3))
v = np.random.randint(-299, 300, size=(32, 98, 3))  # Indexing into `b.shape[2]`; one additional leading dimension for `v`.
result = sliced_sum_numba(b, np.zeros_like(v), v, axis=2)

Это сравнение производительности различных методов, представленных в ответах:

• sliced_reduce
• sliced_sum
• sliced_sum_numba
• reduce_cumulative (первоначальная идея здесь)
• baseline — «классический» цикл Python for (см. ниже).

Примечания по производительности

• sliced_reduce меняет порядок пар индексов с восходящего на убывающий, чтобы превратить вычисление лишних элементов в нулевые; таким образом, однако, массив не проходится в порядке расположения памяти и, кажется, замедляет метод примерно на 30%.
• reduce_cumulative выполняет ряд ненужных операций добавления, которые зависят от распределения начальных и конечных индексов. Для примера OP, где все начальные индексы равны нулю, а конечные индексы распределены равномерно, это будет в среднем в два раза больше операций, чем строго необходимо. Для других распределений (например, с ненулевыми начальными индексами) эта доля вполне может измениться и, следовательно, ухудшить производительность по сравнению с другими методами. Пожалуйста, проверьте свой случай.
• [Отказ от ответственности] Как и в случае со всеми оценками производительности, это приблизительные рекомендации для общего обзора, но они не спасают вас от самостоятельного выполнения тестов производительности для вашего конкретного варианта использования на конкретной машине, чтобы быть абсолютно уверенным в выборе наилучшего варианта.

Используя примеры размеров из ОП:

In [15]: np.random.seed(0)

In [16]: b = np.random.randint(0, 1000, size=(98, 3, 300, 3))

In [17]: v = np.random.randint(-299, 300, size=(32, 98, 3))

In [18]: %timeit sliced_reduce(b, np.zeros_like(v), v, np.add, axis=2)
11.3 ms ± 110 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [19]: %timeit sliced_sum(b, np.zeros_like(v), v, axis=2)
54.9 ms ± 153 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

In [20]: %timeit sliced_sum_numba(b, np.zeros_like(v), v, 2)
16.3 ms ± 609 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [21]: %timeit reduce_cumulative(b, np.zeros_like(v), v, np.add, axis=2)
2.05 ms ± 30.2 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [22]: %timeit baseline(b, np.zeros_like(v), v, axis=2)
79 ms ± 625 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Базовая реализация:

def baseline(a, i, j, axis=None):
    if axis is None:
        axis = len(i.shape)
    i = (a.shape[axis] + i) % a.shape[axis]
    j = (a.shape[axis] + j) % a.shape[axis]
    m = len(i.shape) - axis
    result = np.empty(i.shape + a.shape[axis+1:], dtype=a.dtype)
    for k in np.ndindex(i.shape):
        result[k] = np.sum(a[k[m:] + (slice(i[k], j[k]),)], axis=0)
    return result

Графики производительности

Помимо времени для конкретного примера OP, полезно проверить, как алгоритмы масштабируются с размером массивов данных и индексов. Здесь мы можем разделить фигуры на три разных компонента:

• Начальные измерения индексного массива (те, которых нет в массиве данных). В примере OP это (32,).
• Общие размерности индекса и массива данных (размеры после ведущих до уменьшенной оси). В примере OP это (98, 3).
• Размер оси нужно уменьшить. В примере OP это 300.
• (Конечные размеры массива данных обрабатываются одинаково всеми алгоритмами, и, следовательно, не следует ожидать особого масштабирования.)

Следовательно, мы можем создать графики производительности для трех разных случаев: изменение размера ведущего измерения (размеров), общих размеров и размера оси, которую необходимо уменьшить. Границы выбираются от 1 до N, где N — наибольшая степень 2, при которой ни один задействованный массив не содержит более 5 000 000 элементов (вход, индекс, вывод; промежуточные массивы могут быть больше (например, для sliced_reduce)).

Код см. ниже.

Основные размеры

Общие размеры

Уменьшенный размер

Код

from string import ascii_lowercase as symbols
import numba
import numpy as np
import perfplot

np.random.seed(0)

def sliced_reduce(a, i, j, ufunc=np.add, axis=2):
    indices = np.tile(
        np.repeat(
            np.arange(np.prod(a.shape[:axis])) * a.shape[axis],
            2
        ),
        np.prod(i.shape[:len(i.shape) - axis])
    )
    indices[::2] += (a.shape[axis] + i.ravel()) % a.shape[axis]
    indices[1::2] += (a.shape[axis] + j.ravel()) % a.shape[axis]
    indices = indices.reshape(-1, 2)[::-1].ravel()  # This seems to be counter-effective, please check for your own case.
    result = ufunc.reduceat(a.reshape(-1, *a.shape[axis+1:]), indices)[::2]  # Select only even to odd.
    result[indices[::2] == indices[1::2]] = ufunc.reduce([])
    return result[::-1].reshape(*(i.shape + a.shape[axis+1:]))

def sliced_sum(a, i, j, axis=2):
    l = len(i.shape) - axis
    m = len(i.shape) - l
    n = len(a.shape) - axis - 1
    leading = symbols[:l]
    common = symbols[l:l+m]
    summation = symbols[l+m]
    trailing = symbols[l+m+1:l+m+1+n]
    i = (a.shape[axis] + i) % a.shape[axis]
    j = (a.shape[axis] + j) % a.shape[axis]
    indices, i, j = np.broadcast_arrays(np.arange(a.shape[axis]),
                                        np.expand_dims(i, -1), np.expand_dims(j, -1))
    active_elements = (i <= indices) & (indices < j)
    return np.einsum(f'{leading + common + summation},{common + summation + trailing}->{leading + common + trailing}',
                     active_elements, a)

def sliced_sum_numba(a, i, j, axis=2):
    i = (a.shape[axis] + i) % a.shape[axis]
    j = (a.shape[axis] + j) % a.shape[axis]
    m = np.prod(i.shape[:len(i.shape) - axis], dtype=int)
    n = np.prod(i.shape[len(i.shape) - axis:], dtype=int)
    a_flat = a.reshape(-1, *a.shape[axis:])
    i_flat = i.ravel()
    j_flat = j.ravel()
    result = np.empty((m*n,) + a.shape[axis+1:], dtype=a.dtype)
    numba_sum(a_flat, i_flat, j_flat, m, n, result)
    return result.reshape(*(i.shape + a.shape[axis+1:]))

@numba.jit(parallel=True, nopython=True)
def numba_sum(a, i, j, m, n, out):
    for index in numba.prange(m*n):
        out[index] = np.sum(a[index % n, i[index]:j[index]], axis=0)

def reduce_cumulative(a, i, j, ufunc=np.add, axis=2):
    i = (a.shape[axis] + i) % a.shape[axis]
    j = (a.shape[axis] + j) % a.shape[axis]
    a = np.insert(a, 0, 0, axis)
    c = ufunc.accumulate(a, axis=axis)
    pre = np.ix_(*(range(x) for x in i.shape))
    l = len(i.shape) - axis
    return c[pre[l:] + (j,)] - c[pre[l:] + (i,)]

def baseline(a, i, j, axis=2):
    i = (a.shape[axis] + i) % a.shape[axis]
    j = (a.shape[axis] + j) % a.shape[axis]
    m = len(i.shape) - axis
    result = np.empty(i.shape + a.shape[axis+1:], dtype=a.dtype)
    for k in np.ndindex(i.shape):
        result[k] = np.sum(a[k[m:] + (slice(i[k], j[k]),)], axis=0)
    return result

a = np.random.randint(0, 1000, size=(98, 3, 300, 3))
j = np.random.randint(-299, 300, size=(32, 98, 3))
i = np.zeros_like(j)
check = [f(a, i, j) for f in [sliced_reduce, sliced_sum, sliced_sum_numba, reduce_cumulative, baseline]]
assert all(np.array_equal(check[0], x) for x in check[1:])

perfplot.show(
    # Leading dimensions:
    # setup = lambda n: (np.random.randint(0, 1000, size=(98, 3, 300, 3)),
    #                    np.zeros((n, 98, 3), dtype=int),
    #                    np.random.randint(-299, 300, size=(n, 98, 3))),
    # Common dimensions:
    # setup = lambda n: (np.random.randint(0, 1000, size=(n, 3, 300, 3)),
    #                    np.zeros((32, n, 3), dtype=int),
    #                    np.random.randint(-299, 300, size=(32, n, 3))),
    # Reduced dimension:
    setup = lambda n: (np.random.randint(0, 1000, size=(98, 3, n, 3)),
                       np.zeros((32, 98, 3), dtype=int),
                       np.random.randint(-n+1, n, size=(32, 98, 3))),
    kernels=[
        lambda a: sliced_reduce(*a),
        lambda a: sliced_sum(*a),
        lambda a: sliced_sum_numba(*a),
        lambda a: reduce_cumulative(*a),
        lambda a: baseline(*a),
    ],
    labels=['sliced_reduce', 'sliced_sum', 'sliced_sum_numba', 'reduce_cumulative', 'baseline'],
    # n_range=[2 ** k for k in range(13)],  # Leading dimensions.
    # n_range=[2 ** k for k in range(11)],  # Common dimensions.
    n_range=[2 ** k for k in range(2, 13)],  # Reduced dimension.
    # xlabel='Size of leading dimension',
    # xlabel='Size of first common dimension (second is 3)',
    xlabel='Size of reduced dimension',
)

Другая идея, выдвинутая этот ответ (отсюда вики сообщества), состоит в том, чтобы использовать np.cumsum, а затем выбирать строки, соответствующие индексам слайсов. С нулевыми индексами можно справиться, вставив дополнительную нулевую строку в начале оси, подлежащей сокращению. Этот подход выполняет ненужные вычисления, поскольку он вычисляет полную кумулятивную сумму за пределами конечного индекса. В случае, если стоп-индексы равномерно распределены по оси (с медианой input_array.shape[axis]//2), это в среднем будет выполнять в два раза больше операций сложения, чем необходимо. Тем не менее, этот подход кажется работать довольно хорошо по сравнению с другими методами (по крайней мере, для размеров, указанных в OP).

def reduce_cumulative(a, i, j, ufunc, axis=None):
    if axis is None:
        axis = len(i.shape)
    i = (a.shape[axis] + i) % a.shape[axis]
    j = (a.shape[axis] + j) % a.shape[axis]
    a = np.insert(a, 0, 0, axis)  # Insert zeros to account for zero indices.
    c = ufunc.accumulate(a, axis=axis)
    pre = np.ix_(*(range(x) for x in i.shape))  # Indices for dimensions prior to `axis`.
    l = len(i.shape) - axis  # Number of leading dimensions in `i` and `j`.
    return c[pre[l:] + (j,)] - c[pre[l:] + (i,)]