def METHOD_NAME(self):
    x = np.random.randn(3, 4)
    t = jnp.array(x)
    self.assertTrue(np.allclose(transpose(x), np.asarray(transpose(t))))
    x = np.random.randn(3, 4, 5)
    t = jnp.array(x)
    self.assertTrue(np.allclose(transpose(x, axes=(1, 2, 0)), np.asarray(transpose(t, axes=(1, 2, 0)))))