def body_fn(i, state):
    (rng, x) = state
    (key, subkey) = random.split(rng)
    return (key, f(subkey, x))
def body_fn(i, state):
    (rng, x) = state
    (key, subkey) = random.split(rng)
    return (key, f(subkey, x))