-
Notifications
You must be signed in to change notification settings - Fork 74.1k
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
Refactoring tpu_api_dlsym_initializer #55413
Conversation
@skye could you please review this PR? |
Marking "ready to pull" so I can run internal tests, not actually ready for submission. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This looks good! Just one small comment, and I wanna check the internal tests before approving.
Also, can you test that this doesn't break TF? Since we don't have Cloud TPU presubmits (I think?).
@sshahrokhi Can you please address Ubuntu Sanity errors? Thank you! |
@gbaned the build seems to fail because of a mismatch in bazel version, I guess can we run them again to see if that is resolved? |
@skye is out this week, @michaelbanfield could you please review this? |
PiperOrigin-RevId: 439970863
Seems auto-merge is not happening but the changes are merged into master now, so we can close this. Thank you for the PR. |
We want to remove the static initializer of the TPU for JAX, which stays in the tpu_api_dlsym_initializer file. However, The functions and links in this file are needed to be called and linked in the correct places if we remove the static initializer. Therefore, we are only keeping the static initializer in this file, and have the tpu_initializer_helper file to have the rest of the functions and links. So there is no need to link this file in any place that does not need the static initializer. This PR does not change the behavior of anything, and is only refactoring functions.