TLDR: how to setup nditer when my algorithm needs different number of values for each operand, but I want broadcasting to be applied over the "other" axes.
I'm in the process of converting some algorithms to cython since I've got a lot of looping overhead in the implementation. Originally I implemented the algorithms with support for broadcasting to allow various use-cases, and I would like to keep that.
The algorithms are quite involved, but the issue can be summarized by the following example code:
a = np.arange(5)
b = np.arange(11)
c = 0
for idx in range(len(a)):
c += a[idx] * b[2 * idx]
c += a[idx] * b[2 * idx + 1]
Broadcasting of this could be implemented along the first axis with the exact same code:
a = np.arange(5 * 7).reshape((5, 7, 1))
b = np.arange(11 * 6).reshape((11, 1, 6))
c = 0
for idx in range(a.shape[0]):
c += a[idx] * b[2 * idx]
c += a[idx] * b[2 * idx + 1]
or along the last axis with some slight modifications (not the same result, but that's not of importance here):
a = np.arange(5 * 7).reshape((7, 1, 5))
b = np.arange(11 * 6).reshape((1, 6, 11))
c = 0
for idx in range(a.shape[-1]):
c += a[..., idx] * b[..., 2 * idx]
c += a[..., idx] * b[..., 2 * idx + 1]
The actual algorithms I have need multiple nested for-loops for each "broadcastable unit", can involve more "inputs" (a and b here) and the "output" (c) can also be another array instead of a single value.
When the inner loops are moved over to cython broadcasting is no longer an option. It seems like nditer would be the way to go here, but I cannot figure out how to make it ignore the fact that one of the axes is not broadcastable. I expected that
a = np.arange(5 * 7).reshape((7, 1, 5))
b = np.arange(11 * 6).reshape((1, 6, 11))
it = np.nditer([a, b], flags=['external_loop'])
would allow me to loop over all axes except the one where I apply my custom algorithm, but that does not seem to be the case. Instead I'm met with a ValueError: operands could not be broadcast together with shapes (7,1,5) (1,6,11).
Ideally I would be able to loop as
for a_inner, b_inner, out_inner in it:
out_inner[...] = call_to_cythonized_algorithm(a_inner, b_inner)
where the shapes of the _inner variables match with what I need for the algorithm (5, 11, 0 in the examples above), or potentially with one extra dimension which I could loop over in the cython code.
I've tried a couple other flags as well, but I don't really know what I'm doing, and none of them give me an iterator that works.
Is this possible with the current API, or have I found a limitation of nditer?