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

enable fori-to-scan lowering #6549

Merged
merged 1 commit into from
Apr 23, 2021
Merged

enable fori-to-scan lowering #6549

merged 1 commit into from
Apr 23, 2021

Conversation

mattjj
Copy link
Collaborator

@mattjj mattjj commented Apr 23, 2021

When a lax.fori_loop has static bounds, we can lower it to a scan and thus support efficient reverse-mode autodiff.

We had attempted to add this in #2414, but had to disable it for some weird custom_transforms reason. That means #2026 unblocked it (and maybe #3370 helped too)!.

@google-cla google-cla bot added the cla: yes label Apr 23, 2021
@mattjj mattjj requested a review from hawkinsp April 23, 2021 15:34
@mblondel
Copy link

Awesome, thanks @mattjj! The docstring needs to updated to mention the lowering to scan.

@mattjj mattjj force-pushed the lower-fori-to-scan-2 branch from 1f4ddb9 to 4ac8937 Compare April 23, 2021 15:49
@mattjj mattjj added the pull ready Ready for copybara import and testing label Apr 23, 2021
@copybara-service copybara-service bot merged commit 6969062 into master Apr 23, 2021
@mattjj mattjj deleted the lower-fori-to-scan-2 branch April 23, 2021 17:17
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
cla: yes pull ready Ready for copybara import and testing
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants