[Python] numpy apply_along_axis and function that returns Optional

Trying to understand behavior in apply_along_axis.

The following code will break with an error:

import numpy as np

def func(a):
if a % 3 == 0: return 1
if a % 3 == 1: return 0
else: return None

arr = np.array([[1], [2], [3], [4], [5], [6]])

print(np.apply_along_axis(func, 1, arr))


Error:

TypeError: int() argument must be a string, a bytes-like object or a real number, not 'NoneType'


This is because the numpy code creates a buffer as an array based on the first result of the function, which is dtype=int64 in this case, and can't handle saving a None to the array.

Now, if the very first result from the function is None, numpy will create a buffer of dtype=object instead, and I get a valid result:

arr = np.array([[1], [2], [3], [4], [5], [6]])
# prints [None 1 0 None 1]


Is this expected behavior -- does apply_along_axis just not work when the invoked Python function returns Optional[int]?

Continue reading...