def testReductionOfOutOfBoundsAxis(self):
    x = jnp.ones((3, 4))
    self.assertRaises(ValueError, (lambda : jnp.sum(x, axis=2)))
def testReductionOfOutOfBoundsAxis(self):
    x = jnp.ones((3, 4))
    self.assertRaises(ValueError, (lambda : jnp.sum(x, axis=2)))