feat: infer device_ids and normalize tile assignment#9514

Closed

kvshbg-aws wants to merge 11 commits intopytorch: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
fix for failing ci/cd tests

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

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

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

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

kvshbg-aws
authored and
root
committed
linter fix

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

kvshbg-aws
authored and
root
committed

Commits on Jul 28, 2025

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

Commits on Aug 4, 2025

adding tests, minor bug fixes
kvshbg-aws
committed