Implement `jax.numpy.<fun>.accumulate`

[joeryjoery]

I was trying to run a cumulative minimum when I realized that this isn't part of the main jax.numpy API, while it is the case for numpy.

For now I've managed simply swap it out with:

import jax
import jax.numpy as jnp

key = jax.random.PRNGKey(0)
values = jax.random.normal(key, shape=(10,))

jax.lax.associative_scan(jax.numpy.minimum, values)
>> DeviceArray([-0.3721109, -0.3721109, -0.3721109, -0.7368197, -0.7368197,
             -0.7368197, -0.7368197, -0.7368197, -0.7368197, -0.7368197],            dtype=float32)

# Numpy API:
np.minimum.accumulate(values)
>> array([-0.3721109, -0.3721109, -0.3721109, -0.7368197, -0.7368197,
       -0.7368197, -0.7368197, -0.7368197, -0.7368197, -0.7368197],
      dtype=float32)

So, I have a functional workaround, but this is not as clean as the main numpy API.

Is this perhaps related to #1565?

[jakevdp]

Thanks for the request – I have a prototype of this in #9529; it's unfortunately stalled in review.

[jakevdp]

We're close to getting #9529 marged – ETA probably early next week. On that branch, this works:

In [1]: import jax 
   ...: import jax.numpy as jnp 
   ...:  
   ...: key = jax.random.PRNGKey(0) 
   ...: values = jax.random.normal(key, shape=(10,)) 
   ...:  
   ...: jnp.minimum.accumulate(values)                                                                                                                                                                  
Out[1]: 
DeviceArray([-0.3721109, -0.3721109, -0.3721109, -0.7368197, -0.7368197,
             -0.7368197, -0.7368197, -0.7368197, -0.7368197, -0.7368197],            dtype=float32)

[jakevdp]

Now tracked in #18375