numerical stability and overflow conditions: handling NaN and Inf on GPU

[ianna]

AwkwardArray CPU reducers may return NaN or Inf while the GPU ones would return the zero values:

Here is the full output:

tests-cuda/test_3149_complex_reducers.py ......................F.........ssss                                            [100%]

=========================================================== FAILURES ===========================================================
______________________________________________ test_block_boundary_prod_complex13 ______________________________________________

    def test_block_boundary_prod_complex13():
        np.random.seed(42)
        array = np.random.random(1000)
        complex_array = np.vectorize(complex)(
            array[0 : len(array) : 2], array[1 : len(array) : 2]
        )
        content = ak.contents.NumpyArray(complex_array)
        cuda_content = ak.to_backend(content, "cuda", highlevel=False)
        cpt.assert_allclose(
            ak.prod(cuda_content, -1, highlevel=False),
            ak.prod(content, -1, highlevel=False),
        )
    
        offsets = ak.index.Index64(np.array([0, 5, 996, 1000], dtype=np.int64))
        depth1 = ak.contents.ListOffsetArray(offsets, content)
        cuda_depth1 = ak.to_backend(depth1, "cuda", highlevel=False)
>       cpt.assert_allclose(
            to_list(ak.prod(cuda_depth1, -1, highlevel=False)),
            to_list(ak.prod(depth1, -1, highlevel=False)),
        )

array      = array([0.37454012, 0.95071431, 0.73199394, 0.59865848, 0.15601864,
       0.15599452, 0.05808361, 0.86617615, 0.601115...88, 0.69469647, 0.27214514, 0.59023067, 0.3609739 ,
       0.09158207, 0.91731358, 0.13681863, 0.95023735, 0.44600577])
complex_array = array([0.37454012+0.95071431j, 0.73199394+0.59865848j,
       0.15601864+0.15599452j, 0.05808361+0.86617615j,
       0...469647j,
       0.27214514+0.59023067j, 0.3609739 +0.09158207j,
       0.91731358+0.13681863j, 0.95023735+0.44600577j])
content    = <NumpyArray dtype='complex128' len='500'>
    [0.37454012+0.95071431j 0.73199394+0.59865848j 0.15601864+0.15599452j
  ...469647j 0.27214514+0.59023067j 0.3609739 +0.09158207j
     0.91731358+0.13681863j 0.95023735+0.44600577j]
</NumpyArray>
cuda_content = <NumpyArray dtype='complex128' len='500'>
    [0.37454012+0.95071431j 0.73199394+0.59865848j 0.15601864+0.15599452j
  ...469647j 0.27214514+0.59023067j 0.3609739 +0.09158207j
     0.91731358+0.13681863j 0.95023735+0.44600577j]
</NumpyArray>
cuda_depth1 = <ListOffsetArray len='3'>
    <offsets><Index dtype='int64' len='4'>
        [   0    5  996 1000]
    </Index></offse...739 +0.09158207j
         0.91731358+0.13681863j 0.95023735+0.44600577j]
    </NumpyArray></content>
</ListOffsetArray>
depth1     = <ListOffsetArray len='3'>
    <offsets><Index dtype='int64' len='4'>
        [   0    5  996 1000]
    </Index></offse...739 +0.09158207j
         0.91731358+0.13681863j 0.95023735+0.44600577j]
    </NumpyArray></content>
</ListOffsetArray>
offsets    = <Index dtype='int64' len='4'>
    [   0    5  996 1000]
</Index>

tests-cuda/test_3149_complex_reducers.py:575: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
../../../anaconda3/envs/awkward-py3.11/lib/python3.11/site-packages/cupy/testing/_array.py:24: in assert_allclose
    numpy.testing.assert_allclose(
        actual     = [(0.054843337928496036-0.16291323500570526j), 0j, 0j]
        atol       = 0
        desired    = [(0.05484333792849601-0.16291323500570523j), (nan+nanj), 0j]
        err_msg    = ''
        rtol       = 1e-07
        verbose    = True
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

args = (<function assert_allclose.<locals>.compare at 0x7cc2aba2c860>, array([0.05484334-0.16291324j, 0.        +0.j        ,
       0.        +0.j        ]), array([0.05484334-0.16291324j,        nan       +nanj,
       0.        +0.j        ]))
kwds = {'equal_nan': True, 'err_msg': '', 'header': 'Not equal to tolerance rtol=1e-07, atol=0', 'verbose': True}

    @wraps(func)
    def inner(*args, **kwds):
        with self._recreate_cm():
>           return func(*args, **kwds)
E           AssertionError: 
E           Not equal to tolerance rtol=1e-07, atol=0
E           
E           x and y nan location mismatch:
E            x: array([0.054843-0.162913j, 0.      +0.j      , 0.      +0.j      ])
E            y: array([0.054843-0.162913j,      nan     +nanj, 0.      +0.j      ])

args       = (<function assert_allclose.<locals>.compare at 0x7cc2aba2c860>, array([0.05484334-0.16291324j, 0.        +0.j        ,
       0.        +0.j        ]), array([0.05484334-0.16291324j,        nan       +nanj,
       0.        +0.j        ]))
func       = <function assert_array_compare at 0x7cc2ac8bcf40>
kwds       = {'equal_nan': True, 'err_msg': '', 'header': 'Not equal to tolerance rtol=1e-07, atol=0', 'verbose': True}
self       = <contextlib._GeneratorContextManager object at 0x7cc2ac8d1ed0>

../../../anaconda3/envs/awkward-py3.11/lib/python3.11/contextlib.py:81: AssertionError
=================================================== short test summary info ====================================================
SKIPPED [1] tests-cuda/test_3149_complex_reducers.py:773: awkward_reduce_argmax_complex is not implemented
SKIPPED [1] tests-cuda/test_3149_complex_reducers.py:795: awkward_reduce_argmax_complex is not implemented
SKIPPED [1] tests-cuda/test_3149_complex_reducers.py:817: awkward_reduce_argmin_complex is not implemented
SKIPPED [1] tests-cuda/test_3149_complex_reducers.py:839: awkward_reduce_argmin_complex is not implemented
FAILED tests-cuda/test_3149_complex_reducers.py::test_block_boundary_prod_complex13 - AssertionError: 
=========================================== 1 failed, 31 passed, 4 skipped in 2.85s ============================================

Originally posted by @ianna in #3149 (comment)