def testIssue804(self):
    num_devices = xla_bridge.device_count()
    f = partial(lax.scan, (lambda c, x: ((c + lax.psum(x, 'i')), c)), 0.0)
    api.pmap(f, axis_name='i')(jnp.ones((num_devices, 4)))
def testIssue804(self):
    num_devices = xla_bridge.device_count()
    f = partial(lax.scan, (lambda c, x: ((c + lax.psum(x, 'i')), c)), 0.0)
    api.pmap(f, axis_name='i')(jnp.ones((num_devices, 4)))