Skip to content

[DRAFT] call_jax sets the JAX contextual mesh #9043

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Draft
wants to merge 3 commits into
base: master
Choose a base branch
from
Draft

[DRAFT] call_jax sets the JAX contextual mesh #9043

wants to merge 3 commits into from

Conversation

tengyifei
Copy link
Collaborator

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

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

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

@zpcore zpcore Awaiting requested review from zpcore zpcore will be requested when the pull request is marked ready for review zpcore 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