Skip to content

[while_loop] prepare for torch while_loop signature change.#6872

Merged
JackCaoG merged 1 commit intopytorch:masterfrom
ydwu4:while_loop_sig
Apr 2, 2024
Merged

[while_loop] prepare for torch while_loop signature change.#6872
JackCaoG merged 1 commit intopytorch:masterfrom
ydwu4:while_loop_sig

Conversation

@ydwu4
Copy link
Copy Markdown
Contributor

@ydwu4 ydwu4 commented Apr 2, 2024

In pytorch/pytorch#123018, we add support for closures, params and buffers to while_loop in pytorch but it's failing XLA tests here https://hud.pytorch.org/pr/pytorch/pytorch/123018#23324185080 because the signature of the implantation of DispatchKey.xla needs to be modified. This pr prepares xla key implementation to accept additional_inputs.

Right now, it's just a None value. We need to add proper support for additional_inputs in _xla_while_loop and probably add more tests for closures, parameters and buffers when pytorch/pytorch#123018 ships.

@qihqi
Copy link
Copy Markdown
Collaborator

qihqi commented Apr 2, 2024

cc. @ManfeiBai

Copy link
Copy Markdown
Collaborator

@ManfeiBai ManfeiBai left a comment

Choose a reason for hiding this comment

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

Thanks, LGTM

@ManfeiBai ManfeiBai requested review from JackCaoG, miladm and yeounoh April 2, 2024 17:39
@ydwu4
Copy link
Copy Markdown
Contributor Author

ydwu4 commented Apr 2, 2024

Hi, the tests seem passing. Can we merge it?

@JackCaoG JackCaoG merged commit 02f2a0d into pytorch:master Apr 2, 2024
pytorchmergebot pushed a commit to pytorch/pytorch that referenced this pull request Apr 3, 2024
#123018 introduces a necessary bc breaking change and sees a bunch of xla test failures on CI. We made a pr to pytorch/xla to prepare for the breaking change pytorch/xla#6872. We update the pin of pytorch/xla to reflect the change in this PR.

Pull Request resolved: #123217
Approved by: https://github.com/clee2000
sanketpurandare pushed a commit to sanketpurandare/pytorch that referenced this pull request Apr 22, 2024
pytorch#123018 introduces a necessary bc breaking change and sees a bunch of xla test failures on CI. We made a pr to pytorch/xla to prepare for the breaking change pytorch/xla#6872. We update the pin of pytorch/xla to reflect the change in this PR.

Pull Request resolved: pytorch#123217
Approved by: https://github.com/clee2000
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.

4 participants