@parameterized.named_parameters(jtu.cases_from_list(({'testcase_name': '_shape={}_mode={}_rpadwidth={}_rconstantvalues={}'.format(jtu.format_shape_dtype_string(shape, dtype), mode, pad_width_rank, constant_values_rank), 'shape': shape, 'dtype': dtype, 'mode': mode, 'pad_width_rank': pad_width_rank, 'constant_values_rank': constant_values_rank, 'rng_factory': jtu.rand_default, 'irng_factory': partial(jtu.rand_int, 3)} for (mode, constant_values_rank, shapes) in [('constant', 0, all_shapes), ('constant', 1, all_shapes), ('constant', 2, all_shapes), ('symmetric', None, nonempty_shapes), ('reflect', None, nonempty_shapes), ('wrap', None, nonempty_shapes), ('edge', None, nonempty_shapes)] for (shape, dtype) in _shape_and_dtypes(shapes, all_dtypes) for pad_width_rank in range(3))))
def testPad(self, shape, dtype, mode, pad_width_rank, constant_values_rank, rng_factory, irng_factory):
    rng = rng_factory()
    irng = irng_factory()
    pad_width = irng([len(shape), 2][(2 - pad_width_rank):], onp.int32)

    def onp_fun(x, kwargs):
        if (pad_width.size == 0):
            return x
        return onp.pad(x, pad_width, mode=mode, **kwargs)

    def jnp_fun(x, kwargs):
        return jnp.pad(x, pad_width, mode=mode, **kwargs)

    def args_maker():
        kwargs = {}
        if constant_values_rank:
            kwargs['constant_values'] = rng([len(shape), 2][(2 - constant_values_rank):], dtype)
        return (rng(shape, dtype), kwargs)
    self._CheckAgainstNumpy(onp_fun, jnp_fun, args_maker, check_dtypes=(shape is not jtu.PYTHON_SCALAR_SHAPE))
    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True)
@parameterized.named_parameters(jtu.cases_from_list(({'testcase_name': '_shape={}_mode={}_rpadwidth={}_rconstantvalues={}'.format(jtu.format_shape_dtype_string(shape, dtype), mode, pad_width_rank, constant_values_rank), 'shape': shape, 'dtype': dtype, 'mode': mode, 'pad_width_rank': pad_width_rank, 'constant_values_rank': constant_values_rank, 'rng_factory': jtu.rand_default, 'irng_factory': partial(jtu.rand_int, 3)} for (mode, constant_values_rank, shapes) in [('constant', 0, all_shapes), ('constant', 1, all_shapes), ('constant', 2, all_shapes), ('symmetric', None, nonempty_shapes), ('reflect', None, nonempty_shapes), ('wrap', None, nonempty_shapes), ('edge', None, nonempty_shapes)] for (shape, dtype) in _shape_and_dtypes(shapes, all_dtypes) for pad_width_rank in range(3))))
def testPad(self, shape, dtype, mode, pad_width_rank, constant_values_rank, rng_factory, irng_factory):
    rng = rng_factory()
    irng = irng_factory()
    pad_width = irng([len(shape), 2][(2 - pad_width_rank):], onp.int32)

    def onp_fun(x, kwargs):
        if (pad_width.size == 0):
            return x
        return onp.pad(x, pad_width, mode=mode, **kwargs)

    def jnp_fun(x, kwargs):
        return jnp.pad(x, pad_width, mode=mode, **kwargs)

    def args_maker():
        kwargs = {}
        if constant_values_rank:
            kwargs['constant_values'] = rng([len(shape), 2][(2 - constant_values_rank):], dtype)
        return (rng(shape, dtype), kwargs)
    self._CheckAgainstNumpy(onp_fun, jnp_fun, args_maker, check_dtypes=(shape is not jtu.PYTHON_SCALAR_SHAPE))
    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True)