Skip to content

[DRAFT] call_jax sets the JAX contextual mesh#9043

Draft
tengyifei wants to merge 3 commits intomasterfrom
yifeit/call-jax-mesh-2
Draft

[DRAFT] call_jax sets the JAX contextual mesh#9043
tengyifei wants to merge 3 commits intomasterfrom
yifeit/call-jax-mesh-2

Conversation

@tengyifei
Copy link
Copy Markdown
Collaborator

@tengyifei tengyifei commented Apr 25, 2025

No description provided.

@tengyifei
[call_jax] Bridge the torch_xla and JAX mesh
b8d8be5 
Now we can run a JAX SPMD function that accesses the ambient SPMD mesh
from xb.call_jax.

Fixes #8972.

Also I beefed up the assume_pure tests and updated the docs to mention
that mark_sharding is supported thanks to qihqi@'
#8989.
@tengyifei
update yapf
de22e83
@tengyifei
Address comments
d83ffca
This was referenced Apr 25, 2025
Open
Merged
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@mikegre-google mikegre-google Awaiting requested review from mikegre-google mikegre-google will be requested when the pull request is marked ready for review mikegre-google is a code owner

@qihqi qihqi Awaiting requested review from qihqi qihqi will be requested when the pull request is marked ready for review qihqi is a code owner

@zhanyong-wan zhanyong-wan Awaiting requested review from zhanyong-wan zhanyong-wan will be requested when the pull request is marked ready for review zhanyong-wan is a code owner

@bhavya01 bhavya01 Awaiting requested review from bhavya01 bhavya01 will be requested when the pull request is marked ready for review bhavya01 is a code owner

At least 1 approving review is required to merge this pull request.

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@tengyifei