xla
Use shard_as in scan to ensure that inputs and their gradients have the same sharding
#8879
Merged
Overview
Commits
8
Changes
View On GitHub

Use shard_as in scan to ensure that inputs and their gradients have the same sharding #8879

tengyifei merged 8 commits into master from yifeit/scan-shard-as
tengyifei
tengyifei tengyifei marked this pull request as ready for review 1 year ago
tengyifei tengyifei changed the base branch from master to yifeit/call-jax-cache 1 year ago
tengyifei tengyifei force pushed from 3823d66f to 705814f4 1 year ago
tengyifei tengyifei force pushed from 705814f4 to b2140fa9 1 year ago
tengyifei tengyifei changed the base branch from yifeit/call-jax-cache to master 1 year ago
tengyifei tengyifei force pushed from b2140fa9 to ba4b6ab1 1 year ago
tengyifei tengyifei requested a review from qihqi qihqi 1 year ago
tengyifei tengyifei requested a review from bhavya01 bhavya01 1 year ago
qihqi
qihqi approved these changes on 2025-03-28
tengyifei Use shard_as in scan to ensure that inputs and their gradients have t…
0a85f09e
tengyifei Add back removed API
903119f9
tengyifei Simplify
f625d707
tengyifei Simplify
4637fb09
tengyifei Add test
5c15061e
tengyifei yapf
0bbfd1b8
tengyifei Fix tests
4337f9e7
tengyifei yapf
9dcfe1c3
tengyifei tengyifei force pushed from 5f84ac1d to 9dcfe1c3 1 year ago
tengyifei tengyifei merged 6d88c089 into master 1 year ago

Login to write a write a comment.

Login via GitHub

Reviewers
Assignees
No one assigned
Labels
None yet
Milestone
No milestone
Terms Of Use Privacy Policy Refund Policy Impressum