How to use the numbagg.nanargmax function in numbagg

funcs_reference_funcs = {
    numbagg.allnan: lambda x, **kwargs: np.all(np.isnan(x), **kwargs),
    numbagg.anynan: lambda x, **kwargs: np.any(np.isnan(x), **kwargs),
    numbagg.nansum: np.nansum,
    numbagg.nanmean: np.nanmean,
    numbagg.nanstd: np.nanstd,
    numbagg.nanvar: np.nanvar,
    numbagg.nanmax: np.nanmax,
    numbagg.nanmin: np.nanmin,
    numbagg.count: lambda x, **kwargs: np.sum(~np.isnan(x), **kwargs),
}

argmax_reference_funcs = {
    numbagg.nanargmax: np.nanargmax,
    numbagg.nanargmin: np.nanargmin,
}

moving_references_funcs = {
    numbagg.move_nanmean: wrap_pure_python(numbagg.move_nanmean.func),
}


def allclose(actual, desired, **kwargs):
    if getattr(actual, 'shape', ()) != getattr(desired, 'shape', ()):
        return False
    return np.all(np.isclose(actual, desired, equal_nan=True, **kwargs))


def check_func(numbagg_func, ref_func, x, axis):
    actual = numbagg_func(x, axis=axis)