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

Refactor EXLA Client to use PjRt #423

Merged
merged 23 commits into from
Jun 19, 2021
Merged

Refactor EXLA Client to use PjRt #423

merged 23 commits into from
Jun 19, 2021

Conversation

seanmor5
Copy link
Collaborator

@seanmor5 seanmor5 commented Jun 17, 2021

This PR refactors EXLA to use PjRt. Once merged this will accomplish multiple things:

  • Bumps to stable TF 2.5. Resolves Pin EXLA to a TF Stable Release #62
  • I believe there has been progress upstream on MAC ARM, so it possibly resolves XLA doesn't support Mac ARM #217
  • Makes for far more straightforward asynchronous execution e.g. by exposing BlockHostUntilReady, which will make resolving Introduce Nx.Stream #200 much easier
  • Really reduces the EXLA NIF code footprint so it will be MUCH easier to maintain moving forward
  • Introduces a TPU Client for using Nx with GCP TPU Hosts

As is, fails all async tests on CPU; however, we may want to consider adjusting how async works considering how much simpler this makes things. It also fails two op tests: QR Decomposition and Window Max. I'm investigating whether or not this might be an upstream issue.

Fails most tests on TPU. This is an issue with the executable not liking the layout of our input buffers. I've been talking with some of the PjRt developers about this issue.

Yet to test on GPU, but I don't think there will be any major hiccups.

I will need to adjust our Makefile to try and get rid of the NumPy dependency before this is merged.

@seanmor5 seanmor5 merged commit 9391672 into main Jun 19, 2021
@seanmor5 seanmor5 deleted the sm-tpu branch June 19, 2021 14:01
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

XLA doesn't support Mac ARM Pin EXLA to a TF Stable Release
2 participants