def testMap(self):
    f = (lambda x: (x ** 2))
    xs = jnp.arange(10)
    expected = (xs ** 2)
    actual = lax.map(f, xs)
    self.assertAllClose(actual, expected)
def testMap(self):
    f = (lambda x: (x ** 2))
    xs = jnp.arange(10)
    expected = (xs ** 2)
    actual = lax.map(f, xs)
    self.assertAllClose(actual, expected)