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 requested a review from

qihqi 1 year ago

zzzwen force pushed from 8d8421b8 to 2cd1f803 1 year ago

zzzwen requested a review from

hfan 1 year ago

zzzwen force pushed from c9885f5f to cbb1aafd 1 year ago

zzzwen force pushed from cbb1aafd to d6ae1b16 1 year ago

qihqi approved these changes on 2025-06-19

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 force pushed from 9557761f to 74d15ef8 1 year ago

Clean up comment

5ec61278

zzzwen enabled auto-merge (squash) 1 year ago

zzzwen removed review request from

hfan 1 year ago

lint

312047ea

zzzwen merged d780e55c into master 1 year ago

Reviewers

qihqi

Assignees

No one assigned

Labels

None yet

Milestone

No milestone