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

Add jtu.count_device_put for tests to count device_put. #5998

Merged
merged 3 commits into from
Mar 10, 2021

Conversation

zhangqiaorjc
Copy link
Collaborator

No description provided.

@google-cla google-cla bot added the cla: yes label Mar 9, 2021
@zhangqiaorjc zhangqiaorjc requested a review from mattjj March 9, 2021 22:33
jax/_src/lax/lax.py Outdated Show resolved Hide resolved
@mattjj
Copy link
Collaborator

mattjj commented Mar 10, 2021

I had pushed some commits on top of @zhangqiaorjc's, but those are trickier to land because they involve some changes to jaxprs. So I'm going to try splitting those off into a separate PR!

@copybara-service copybara-service bot merged commit cf9b77f into jax-ml:master Mar 10, 2021
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