xla
Add jax_device context manager to control the device target
#9382
Merged

Add jax_device context manager to control the device target #9382

zzzwen merged 7 commits into pytorch:master from zzzwen:to_device
zzzwen
zzzwen zzzwen requested a review from qihqi qihqi 301 days ago
zzzwen zzzwen force pushed from 8d8421b8 to 2cd1f803 301 days ago
zzzwen zzzwen requested a review from hfan hfan 301 days ago
zzzwen zzzwen force pushed from c9885f5f to cbb1aafd 301 days ago
zzzwen zzzwen force pushed from cbb1aafd to d6ae1b16 301 days ago
qihqi
qihqi approved these changes on 2025-06-19
qihqi
qihqi commented on 2025-06-19
Add jax_device context manager to control the device target
5cb53433
Add some comment
ea7c1d14
Add default to match
39a99804
lint
279c06ad
fix test
74d15ef8
zzzwen zzzwen force pushed from 9557761f to 74d15ef8 298 days ago
Clean up comment
5ec61278
zzzwen zzzwen enabled auto-merge (squash) 298 days ago
zzzwen zzzwen removed review request from hfan hfan 298 days ago
lint
312047ea
zzzwen zzzwen merged d780e55c into master 298 days ago

Login to write a write a comment.

Login via GitHub

Reviewers
Assignees
No one assigned
Labels
Milestone