def testMapEmpty(self):
    ans = lax.map((lambda x: (x * x)), jnp.array([]))
    expected = jnp.array([])
    self.assertAllClose(ans, expected)
def testMapEmpty(self):
    ans = lax.map((lambda x: (x * x)), jnp.array([]))
    expected = jnp.array([])
    self.assertAllClose(ans, expected)