def test_single(m, args_maker, repeats, axis):
    lax_ans = jnp.repeat(m, repeats, axis)
    numpy_ans = onp.repeat(m, repeats, axis)
    self.assertAllClose(lax_ans, numpy_ans, check_dtypes=True, rtol=tol, atol=tol)
    jnp_fun = (lambda arg: jnp.repeat(arg, repeats=repeats, axis=axis))
    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True)
def test_single(m, args_maker, repeats, axis):
    lax_ans = jnp.repeat(m, repeats, axis)
    numpy_ans = onp.repeat(m, repeats, axis)
    self.assertAllClose(lax_ans, numpy_ans, check_dtypes=True, rtol=tol, atol=tol)
    jnp_fun = (lambda arg: jnp.repeat(arg, repeats=repeats, axis=axis))
    self._CompileAndCheck(jnp_fun, args_maker, check_dtypes=True)