@parameterized.named_parameters(jtu.cases_from_list(({'testcase_name': '_shape={}_offset={}_axis1={}_axis2={}'.format(jtu.format_shape_dtype_string(shape, dtype), offset, axis1, axis2), 'dtype': dtype, 'shape': shape, 'offset': offset, 'axis1': axis1, 'axis2': axis2, 'rng_factory': jtu.rand_default} for dtype in default_dtypes for shape in [shape for shape in all_shapes if (len(shape) >= 2)] for axis1 in range((- len(shape)), len(shape)) for axis2 in [a for a in range((- len(shape)), len(shape)) if ((a % len(shape)) != (axis1 % len(shape)))] for offset in list(range((- 4), 4)))))
def testDiagonal(self, shape, dtype, offset, axis1, axis2, rng_factory):
    rng = rng_factory()
    onp_fun = (lambda arg: onp.diagonal(arg, offset, axis1, axis2))
    jnp_fun = (lambda arg: jnp.diagonal(arg, offset, axis1, axis2))
    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': '_shape={}_offset={}_axis1={}_axis2={}'.format(jtu.format_shape_dtype_string(shape, dtype), offset, axis1, axis2), 'dtype': dtype, 'shape': shape, 'offset': offset, 'axis1': axis1, 'axis2': axis2, 'rng_factory': jtu.rand_default} for dtype in default_dtypes for shape in [shape for shape in all_shapes if (len(shape) >= 2)] for axis1 in range((- len(shape)), len(shape)) for axis2 in [a for a in range((- len(shape)), len(shape)) if ((a % len(shape)) != (axis1 % len(shape)))] for offset in list(range((- 4), 4)))))
def testDiagonal(self, shape, dtype, offset, axis1, axis2, rng_factory):
    rng = rng_factory()
    onp_fun = (lambda arg: onp.diagonal(arg, offset, axis1, axis2))
    jnp_fun = (lambda arg: jnp.diagonal(arg, offset, axis1, axis2))
    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)