@parameterized.named_parameters(jtu.cases_from_list(({'testcase_name': 'op={}_shape=[{}]_axis={}_out_dtype={}'.format(op, jtu.format_shape_dtype_string(shape, dtype), axis, out_dtype.__name__), 'axis': axis, 'shape': shape, 'dtype': dtype, 'out_dtype': out_dtype, 'rng_factory': jtu.rand_default, 'jnp_op': getattr(jnp, op), 'onp_op': getattr(onp, op)} for op in ['cumsum', 'cumprod'] for dtype in all_dtypes for out_dtype in default_dtypes for shape in all_shapes for axis in ([None] + list(range((- len(shape)), len(shape)))))))
def testCumSumProd(self, axis, shape, dtype, out_dtype, onp_op, jnp_op, rng_factory):
    rng = rng_factory()
    onp_fun = (lambda arg: onp_op(arg, axis=axis, dtype=out_dtype))
    onp_fun = jtu.ignore_warning(category=onp.ComplexWarning)(onp_fun)
    jnp_fun = (lambda arg: jnp_op(arg, axis=axis, dtype=out_dtype))
    jnp_fun = jtu.ignore_warning(category=jnp.ComplexWarning)(jnp_fun)
    args_maker = (lambda : [rng(shape, dtype)])
    tol = max(jtu.tolerance(dtype), jtu.tolerance(out_dtype))
    self._CheckAgainstNumpy(onp_fun, jnp_fun, args_maker, check_dtypes=True, tol=tol)
    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True)
@parameterized.named_parameters(jtu.cases_from_list(({'testcase_name': 'op={}_shape=[{}]_axis={}_out_dtype={}'.format(op, jtu.format_shape_dtype_string(shape, dtype), axis, out_dtype.__name__), 'axis': axis, 'shape': shape, 'dtype': dtype, 'out_dtype': out_dtype, 'rng_factory': jtu.rand_default, 'jnp_op': getattr(jnp, op), 'onp_op': getattr(onp, op)} for op in ['cumsum', 'cumprod'] for dtype in all_dtypes for out_dtype in default_dtypes for shape in all_shapes for axis in ([None] + list(range((- len(shape)), len(shape)))))))
def testCumSumProd(self, axis, shape, dtype, out_dtype, onp_op, jnp_op, rng_factory):
    rng = rng_factory()
    onp_fun = (lambda arg: onp_op(arg, axis=axis, dtype=out_dtype))
    onp_fun = jtu.ignore_warning(category=onp.ComplexWarning)(onp_fun)
    jnp_fun = (lambda arg: jnp_op(arg, axis=axis, dtype=out_dtype))
    jnp_fun = jtu.ignore_warning(category=jnp.ComplexWarning)(jnp_fun)
    args_maker = (lambda : [rng(shape, dtype)])
    tol = max(jtu.tolerance(dtype), jtu.tolerance(out_dtype))
    self._CheckAgainstNumpy(onp_fun, jnp_fun, args_maker, check_dtypes=True, tol=tol)
    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True)