feat: infer device_ids and normalize tile assignment - #9514

Closed

kvshbg-aws wants to merge 11 commits into

pytorch:masterfrom

kvshbg-aws:kvshbg-aws/local-spmd-normalize-infer

Commits on Jul 27, 2025

feat: abstraction of xla::OpSharding proto using wrapper class

kvshbg-aws
authored and
root
committed
c8ae644
fix for failing ci/cd tests

kvshbg-aws
authored and
root
committed
fda80a5
small fix to remove unwanted var

kvshbg-aws
authored and
root
committed
7ae584e
bug fix: set value of denormalized_tile_assignment when sharding.tile_assignment() is empty

kvshbg-aws
authored and
root
committed
cce2ca9
use tensors to get denormalized_tile_assignment directly instead of po_data

kvshbg-aws
authored and
root
committed
1ddfaf4
use tensors instead of paramters_data to get denormalized_tile_assignment

kvshbg-aws
authored and
root
committed
00ded46
linter fix

kvshbg-aws
authored and
root
committed
3c07da6
feat:infer addressable devices for submeshing
root
committed
0e229b5
fix: remove duplicate function

kvshbg-aws
authored and
root
committed
21ad6c4

Commits on Jul 28, 2025

bug fix for argument_handles for non-zero initialized sub-meshes; linter fixes
kvshbg-aws
committed
9fae71e

Commits on Aug 4, 2025

adding tests, minor bug fixes
kvshbg-aws
committed
fc670ef