Skip to content

[runtime:pallas] migrate to torch_tpu's new Pallas buffer donation API#2351

Open
cota wants to merge 2 commits into
pytorch:mainfrom
cota:donation2
Open

[runtime:pallas] migrate to torch_tpu's new Pallas buffer donation API#2351
cota wants to merge 2 commits into
pytorch:mainfrom
cota:donation2

Conversation

@cota
Copy link
Copy Markdown
Collaborator

@cota cota commented May 7, 2026

input_output_aliases will go away. Instead, implement aliasing by using donate_argnums (which is now strictly for donation) and explicitly copying back results via copy_().

cota added 2 commits May 7, 2026 15:15
This is a non-functional change that reduces code duplication.
`input_output_aliases` will go away. Instead, implement aliasing
by using `donate_argnums` (which is now strictly for donation) and
explicitly copying back results via `copy_()`.
@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label May 7, 2026
Copy link
Copy Markdown
Contributor

@oulgen oulgen left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

explicitly copying back results via copy_()

Isn't this very inefficient?

@cota
Copy link
Copy Markdown
Collaborator Author

cota commented May 7, 2026

Yes. As discussed on chat, the idea is to get off the deprecated API first, then optimize.

@cota cota marked this pull request as ready for review May 8, 2026 10:34
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Meta Open Source bot.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants