Searched refs:partial_sharding (Results 1 – 2 of 2) sorted by relevance
304 const HloSharding& partial_sharding, const HloSharding& target_sharding) { in PartialReplicateReshardCompatibleSharding() argument305 if (!partial_sharding.ReplicateOnLastTileDim()) { in PartialReplicateReshardCompatibleSharding()308 int64 rank = partial_sharding.tile_assignment().num_dimensions() - 1; in PartialReplicateReshardCompatibleSharding()316 partial_sharding.tile_assignment().Each( in PartialReplicateReshardCompatibleSharding()320 gid *= partial_sharding.tile_assignment().dim(i); in PartialReplicateReshardCompatibleSharding()334 int64 partial_tile_size = partial_sharding.tile_assignment().dim(dim); in PartialReplicateReshardCompatibleSharding()353 auto reshape_dimensions = partial_sharding.tile_assignment().dimensions(); in PartialReplicateReshardCompatibleSharding()368 auto reshape_tile_assignment = partial_sharding.tile_assignment(); in PartialReplicateReshardCompatibleSharding()1437 bool partial_sharding = false; in UngroupSharding() local1444 partial_sharding = true; in UngroupSharding()[all …]
385 const HloSharding& partial_sharding, const HloSharding& target_sharding);