I got around this with some (masked) numpy operations. perhaps it is useful? I was seeing the np.argmax results on entries with all NaN evaluate to zero, which was not useful since the axis I was computing argmax across had valid entries if the result was 0 (think 0-index month, i.e., January, within a year). So I did this instead:

```

test_arr is some array with some nodata value, and is of dims [channels, rows, columns]

nodata = -32768 ma = np.ma.masked_equal(test_arr, nodata)

use np.any to get a mask of rows/columns which have all masked entries

spec_axis = 0 all_na_mask = np.any(ma, axis=spec_axis)

get the argmax across specified axis

argm = np.argmax(test_arr, axis=spec_axis) argm = np.ma.masked_less(argm, -np.inf) argm.mask = ~all_na_mask ```

big piece here is modifying the mask directly and making sure that is correct. numpy docs advise against this approach but it seems to be giving me what I want.