Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Think about parallel execution of our mapAccums and their derivatives (associative operations?) #109

Open
Mikolaj opened this issue Feb 25, 2024 · 0 comments
Labels
help wanted Extra attention is needed postpone

Comments

@Mikolaj
Copy link
Owner

Mikolaj commented Feb 25, 2024

When the implementation of fold, scan, mapAccum and their derivatives settles down, let's think about parallel execution. A GPU backend to actually benchmark the result wouldn't hurt either. Tom provides

https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.associative_scan.html

At the time of writing, everything is a special case of mapAccumL or mapAccumR, but it may very well make sense to re-introduce a standalone fold just to get a better parallelism (overriding a default method may be a good start) and then we may want to generalize the type of the fold method to any heterogeneous vectors of tensors, which by that time may get easier to use than currently (via typing, adaptors or otherwise). If if turns out scan is easier to parallelise than mapAccum (unlikely?), we may also recover the derivative expressed in terms of scan and build instead of in terms of mapAccum, at the cost of applying the derivative of the little function twice (details in old commits).

@Mikolaj Mikolaj added help wanted Extra attention is needed postpone labels Feb 25, 2024
@Mikolaj Mikolaj changed the title This about parallel execution of our mapAccums and their derivatives (associative operations?) Think about parallel execution of our mapAccums and their derivatives (associative operations?) Feb 26, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
help wanted Extra attention is needed postpone
Projects
None yet
Development

No branches or pull requests

1 participant