如何实现任意维度同形状NumPy数组的自定义元素级二元操作？

针对你需要对任意维度的同形状NumPy数组执行自定义二元元素级操作的需求，这里提供几种实用方案：

方法1：用np.nditer遍历任意维度数组 #

np.nditer是NumPy原生的多维数组迭代工具，能自动适配数组的维度数量，无需手动编写多层循环，同时支持高效的元素访问与赋值。

import numpy as np

def custom_elementwise_op(A, B, op):
    assert A.shape == B.shape, "两个数组的形状必须完全一致"
    result = np.zeros_like(A)
    # 同时迭代A、B的只读元素，以及result的可写元素
    for a_val, b_val, res_val in np.nditer([A, B, result], 
                                          op_flags=[['readonly'], ['readonly'], ['writeonly']]):
        res_val[...] = op(a_val, b_val)
    return result

# 示例自定义二元操作：a*b + a²
def my_custom_op(a, b):
    return a * b + a**2

# 测试3维数组场景
A = np.random.rand(2, 3, 4)
B = np.random.rand(2, 3, 4)
output = custom_elementwise_op(A, B, my_custom_op)
# 验证结果正确性
assert np.allclose(output, A*B + A**2)

方法2：展平数组为一维处理 #

将多维数组展平为一维后处理，再恢复原形状，逻辑简单直观，适合快速实现：

def custom_elementwise_op_flat(A, B, op):
    assert A.shape == B.shape
    A_flat = A.flatten()
    B_flat = B.flatten()
    # 遍历一维数组执行自定义操作
    result_flat = np.array([op(a, b) for a, b in zip(A_flat, B_flat)], dtype=A.dtype)
    return result_flat.reshape(A.shape)

# 验证
output_flat = custom_elementwise_op_flat(A, B, my_custom_op)
assert np.allclose(output_flat, output)

方法3：矢量化自定义函数（推荐高性能场景） #

如果你的自定义操作可以拆分为NumPy的矢量化运算，直接利用广播机制处理，这是效率最高的方案——完全依托NumPy的底层优化，避免Python层面的循环：

# 直接编写矢量化的自定义操作函数
def my_op_vectorized(A, B):
    return A * B + A**2

output_vectorized = my_op_vectorized(A, B)
assert np.allclose(output_vectorized, output)

只要复杂逻辑能通过NumPy的内置函数组合实现，优先选择这种方式。

方法4：用np.vectorize包装标量函数 #

np.vectorize可以将标量级的自定义函数转换为支持数组输入的函数，本质是语法糖（底层仍为循环），适合快速实现但对性能要求不高的场景：

vectorized_op = np.vectorize(my_custom_op)
output_vec = vectorized_op(A, B)
assert np.allclose(output_vec, output)

内容的提问来源于stack exchange，提问作者Tamojit Maiti