Strange behavior of jax.lax.scan

[onnoeberhard]

I am trying to compute a sum as follows:

# Function definitions (m, dy, dx, E, T_rev, mu are all fixed numbers/arrays)
ys = jnp.array(list(itertools.product(*((m + 1) * [np.arange(dy)]))))
@jax.jit
def z(y):
    @jax.vmap
    def f(x):
        e = 0
        x_ = x
        for yi in y:
            e += jnp.log(E[yi] @ x_)
            x_ = T_rev @ x_
        return jnp.exp(e) * (w[*y] - v @ x)**2
    return mu @ f(jnp.eye(dx))

# Sum
s = 0.
for y in ys:
    s += z(y)

In Jax, this can be implemented as

zs = jax.vmap(z)(ys)
s = zs.sum()

However, ys can become very large, so I don't want to use vmap because it will allocate too much memory. Instead, I want to use jax.lax.scan:

c, zs_ = jax.lax.scan(lambda c, y: (c + z(y), z(y)), 0., ys)

Here, zs_ is equal (up to floating-point precision) to zs, but c is very different from s (only accurate up to the first two digits). What is going on here? I thought it must be the rolling summation, but

c_, = jax.lax.scan(lambda c, y: (c + z, None), 0., zs)

returns the correct result (c_ = s). Any help is appreciated! The function z is not stateful, all the other parameters (mu etc. are fixed in advance).

[onnoeberhard]

Okay, it seems that the error is indeed due to the summation of many small floating point numbers. Replacing the naive addition by a Kahanen addition fixes the problem. I can only assume that the last example I show above (jax.lax.scan(lambda c, y: (c + z, None), 0., zs)) works because Jax does some clever compilation that turns it into a non-rolling sum.