def test_while_loop_of_pmap(self):

    def body(i, x):
        result = api.pmap((lambda z: lax.psum(jnp.sin(z), 'i')), axis_name='i')(x)
        return (result + x)
    f_loop = (lambda x: lax.fori_loop(0, 3, body, x))
    ans = f_loop(jnp.ones(api.device_count()))
    del body, f_loop

    def body2(i, x):
        result = jnp.broadcast_to(jnp.sin(x).sum(), x.shape)
        return (result + x)
    g_loop = (lambda x: lax.fori_loop(0, 3, body2, x))
    expected = g_loop(jnp.ones(api.device_count()))
    self.assertAllClose(ans, expected, check_dtypes=False)
def test_while_loop_of_pmap(self):

    def body(i, x):
        result = api.pmap((lambda z: lax.psum(jnp.sin(z), 'i')), axis_name='i')(x)
        return (result + x)
    f_loop = (lambda x: lax.fori_loop(0, 3, body, x))
    ans = f_loop(jnp.ones(api.device_count()))
    del body, f_loop

    def body2(i, x):
        result = jnp.broadcast_to(jnp.sin(x).sum(), x.shape)
        return (result + x)
    g_loop = (lambda x: lax.fori_loop(0, 3, body2, x))
    expected = g_loop(jnp.ones(api.device_count()))
    self.assertAllClose(ans, expected, check_dtypes=False)