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.

Sign up for GitHub

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

Jump to bottom

[nnx] vectorize vmap split counts #3989

Merged

copybara-service merged 1 commit into main from nnx-fix-vmap

Jun 12, 2024

Collaborator

cgarciae commented Jun 12, 2024

What does this PR do?

Vectorizes split RngCount state in vmap, this avoids issues that arise when rows perform a variable number of splits (e.g. using cond).
Removes RngKeyBackup and RngStream.key_backups.


          [nnx] vectorize vmap split counts

e4f417a

cgarciae added the pull ready label

chiamp approved these changes

View reviewed changes

copybara-service bot merged commit 551eaf0 into main

21 checks passed

copybara-service bot deleted the nnx-fix-vmap branch

June 12, 2024 20:13

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels