多维数组任意轴求和归约的泛化实现方法问询

Great question! The key to ditching those hardcoded nested loops and making this work for any dimension or axis is to lean into the stride information and dynamic indexing of the non-reduction axes. Let’s break this down step by step, with code examples to make it concrete.

Core Approach #

Instead of writing nested loops for each fixed dimension, we split the problem into two reusable parts:

1. Iterate over every combination of indices excluding the reduction axis (each combination maps to one element in our output array).
2. For each combination, calculate its starting position in the flat row-major array, then accumulate values along the reduction axis using its stride value.

Step-by-Step Implementation #

Let’s assume we’re working with these inputs (adjust types based on your language/use case):

• data: The flat 1D array storing our multi-dimensional matrix (row-major order).
• sizes: A list where sizes[d] is the length of the d-th dimension (0-based).
• strides: A list where strides[d] is the number of elements to skip in data to move one step along the d-th dimension.
• axis: The 0-based axis we want to sum over.

1. Validate Input & Set Up Output Dimensions

First, make sure the target axis is valid, then define the output array’s dimensions and strides:

def validate_axis(axis, num_dims):
    if not 0 <= axis < num_dims:
        raise ValueError(f"Axis {axis} is out of bounds for a {num_dims}-dimensional array")

num_dims = len(sizes)
validate_axis(axis, num_dims)

# Output dimensions: remove the reduction axis from the original sizes
new_sizes = sizes[:axis] + sizes[axis+1:]
# Output strides: same as original, minus the reduction axis's stride
new_strides = strides[:axis] + strides[axis+1:]

2. Calculate Total Output Elements

Figure out how many elements we need to generate in the reduced array:

from functools import reduce
import operator

total_output = reduce(operator.mul, new_sizes, 1)
result = [0.0] * total_output  # Use int or other types if your data isn't floating-point

3. Map Linear Indices to Multi-Dimensional Indices

We need a way to convert a linear index (like i in a 1D loop) to the corresponding multi-dimensional indices for the output array. Here’s a simple helper function for that:

def unravel_index(linear_idx, sizes):
    idx = []
    remaining = linear_idx
    # Work backwards from the last dimension to the first
    for s in reversed(sizes):
        remaining, rem = divmod(remaining, s)
        idx.append(rem)
    # Reverse to get the correct order of indices
    return list(reversed(idx))

4. Run the Sum Reduction

Loop over each output element, compute its starting offset in the original flat array, then sum all elements along the reduction axis:

axis_size = sizes[axis]
axis_stride = strides[axis]

for i in range(total_output):
    # Get the multi-dimensional indices for this output element
    output_idx = unravel_index(i, new_sizes)
    
    # Calculate the starting position in the original flat array
    start_offset = 0
    for j in range(len(output_idx)):
        # Map the output index to its original dimension (skip the reduction axis)
        original_dim = j if j < axis else j + 1
        start_offset += output_idx[j] * strides[original_dim]
    
    # Accumulate values along the reduction axis
    sum_val = 0.0
    for k in range(axis_size):
        sum_val += data[start_offset + k * axis_stride]
    
    result[i] = sum_val

Why This Works #

• No Hardcoding: The code adapts automatically to any number of dimensions or target axis—no need to rewrite loops when your matrix shape changes.
• Stride-Powered Efficiency: We use the row-major stride values to directly jump to the correct elements in the flat array, avoiding manual calculations of dimension products.
• Universal Compatibility: This logic works for any row-major stored matrix, regardless of how many dimensions it has.

Quick Test Example #

Let’s say we have a 4D array with sizes = [2, 3, 4, 5] (row-major, so strides = [60, 20, 5, 1]). If we reduce along axis 1:

• The output becomes a 3D array with new_sizes = [2, 4, 5]
• For the first output element (i=0), output_idx = [0,0,0], starting offset is 0. We sum data[0], data[20], data[40] (3 elements, matching axis_size=3).
• This gives exactly the same result as hardcoded nested loops, but works for any dimension/axis.

Optimization Tips #

• For large arrays, precompute all starting offsets upfront to avoid recalculating them in the loop.
• In compiled languages like C++, you can inline the unravel_index logic with arithmetic operations to speed things up.
• If your hardware supports it, vectorize the inner sum loop (though note that elements along non-contiguous axes won’t be in sequential memory).

内容的提问来源于stack exchange，提问作者Tom de Geus