@parameterized.named_parameters(jtu.cases_from_list(({'testcase_name': '{}_inshape={}_axis={}'.format(rec.test_name.capitalize(), jtu.format_shape_dtype_string(shape, dtype), axis), 'rng_factory': rec.rng_factory, 'shape': shape, 'dtype': dtype, 'onp_op': getattr(onp, rec.name), 'jnp_op': getattr(jnp, rec.name), 'axis': axis} for rec in JAX_ARGMINMAX_RECORDS for (shape, dtype) in _shape_and_dtypes(rec.shapes, rec.dtypes) for axis in range((- len(shape)), len(shape)))))
def testArgMinMax(self, onp_op, jnp_op, rng_factory, shape, dtype, axis):
    rng = rng_factory()
    if ((dtype == onp.complex128) and (jtu.device_under_test() == 'gpu')):
        raise unittest.SkipTest('complex128 reductions not supported on GPU')

    def onp_fun(array_to_reduce):
        return onp_op(array_to_reduce, axis).astype(jnp.int_)

    def jnp_fun(array_to_reduce):
        return jnp_op(array_to_reduce, axis)
    args_maker = (lambda : [rng(shape, dtype)])
    self._CheckAgainstNumpy(onp_fun, jnp_fun, args_maker, check_dtypes=True)
    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True)
@parameterized.named_parameters(jtu.cases_from_list(({'testcase_name': '{}_inshape={}_axis={}'.format(rec.test_name.capitalize(), jtu.format_shape_dtype_string(shape, dtype), axis), 'rng_factory': rec.rng_factory, 'shape': shape, 'dtype': dtype, 'onp_op': getattr(onp, rec.name), 'jnp_op': getattr(jnp, rec.name), 'axis': axis} for rec in JAX_ARGMINMAX_RECORDS for (shape, dtype) in _shape_and_dtypes(rec.shapes, rec.dtypes) for axis in range((- len(shape)), len(shape)))))
def testArgMinMax(self, onp_op, jnp_op, rng_factory, shape, dtype, axis):
    rng = rng_factory()
    if ((dtype == onp.complex128) and (jtu.device_under_test() == 'gpu')):
        raise unittest.SkipTest('complex128 reductions not supported on GPU')

    def onp_fun(array_to_reduce):
        return onp_op(array_to_reduce, axis).astype(jnp.int_)

    def jnp_fun(array_to_reduce):
        return jnp_op(array_to_reduce, axis)
    args_maker = (lambda : [rng(shape, dtype)])
    self._CheckAgainstNumpy(onp_fun, jnp_fun, args_maker, check_dtypes=True)
    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True)