Add some sort of batched reduction to desc.batching

[f0uriest]

          not for this PR, but at some point might be worth implementing some sort of batched reduction `desc.batching`

Originally posted by @f0uriest in #1440 (comment)

Basically, if we need to evaluate some expensive function over a bunch of inputs and then reduce it (eg sum). Currently the two main options are a single loop which is slow on GPU, or a full vmap which means materializing the full array in memory (or perhaps hoping the the compiler fuses stuff so it isn't but still). One could potentially improve the performance of the loop by unrolling part of it though that increases compile time. Ideally you want something that takes advantage of the fact that the loop iterations are independent.

semantically should be equivalent to the following:

def unbatched_reduce(fun, x, reduction=jnp.add):
    out = 0
    for xi in x:
        out = reduction(out, fun(xi))
    return out