Skip to main content

Browse the archive

Revision 824ac86620572285b86dd09529f8869ef36883ad authored by Peter Hawkins on 03 April 2020, 19:39:56 UTC, committed by GitHub on 03 April 2020, 19:39:56 UTC 
Reimplement np.cumsum and np.cumprod in terms of a parallel prefix scan. (#2596)
 
* Reimplement np.cumsum and np.cumprod in terms of a parallel prefix scan.

Unlike the existing implementation based on lax.reduce_window, this implementation is O(n log n) instead of O(n^2) and is arbitrarily differentiable.

Fixes #1212, #2418, #2542.
May help with issue #2380.

* Relax gradient test tolerance.
1 parent 192e908
History
File Mode Size
benchmark.py -rw-r--r-- 4.0 KB
pmap_benchmark.py -rw-r--r-- 3.8 KB
back to top