xla
[DRAFT] call_jax sets the JAX contextual mesh
#9043

Open

[DRAFT] call_jax sets the JAX contextual mesh #9043

tengyifei wants to merge 3 commits into master from yifeit/call-jax-mesh-2

tengyifei

tengyifei

[call_jax] Bridge the torch_xla and JAX mesh

b8d8be50

tengyifei

update yapf

de22e83b

tengyifei

Address comments

d83ffca0

Login to write a write a comment.

Login via GitHub

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Milestone

No milestone