def cumprod(x):
    s = jnp.ones((2, 32), jnp.float32)
    return lax.scan((lambda s, x: ((x * s), s)), s, x)
def cumprod(x):
    s = jnp.ones((2, 32), jnp.float32)
    return lax.scan((lambda s, x: ((x * s), s)), s, x)