def testTraceMethod(self):
    x = onp.random.randn(3, 4).astype(jnp.float_)
    self.assertAllClose(x.trace(), jnp.array(x).trace(), check_dtypes=True)
    self.assertAllClose(x.trace(), api.jit((lambda y: y.trace()))(x), check_dtypes=True)
def testTraceMethod(self):
    x = onp.random.randn(3, 4).astype(jnp.float_)
    self.assertAllClose(x.trace(), jnp.array(x).trace(), check_dtypes=True)
    self.assertAllClose(x.trace(), api.jit((lambda y: y.trace()))(x), check_dtypes=True)