def scan_reference(f, init, xs):
    carry = init
    ys = []
    for x in xs:
        (carry, y) = f(carry, x)
        ys.append(lax.reshape(y, ((1,) + np.shape(y))))
    ys = lax.concatenate(ys, 0)
    return (carry, ys)
def scan_reference(f, init, xs):
    carry = init
    ys = []
    for x in xs:
        (carry, y) = f(carry, x)
        ys.append(lax.reshape(y, ((1,) + np.shape(y))))
    ys = lax.concatenate(ys, 0)
    return (carry, ys)