def body(i, x):
    result = api.pmap((lambda z: lax.psum(jnp.sin(z), 'i')), axis_name='i')(x)
    return (result + x)
def body(i, x):
    result = api.pmap((lambda z: lax.psum(jnp.sin(z), 'i')), axis_name='i')(x)
    return (result + x)