Searched refs:target_sharding (Results 1 – 2 of 2) sorted by relevance
304 const HloSharding& partial_sharding, const HloSharding& target_sharding) { in PartialReplicateReshardCompatibleSharding() argument309 int64 target_rank = target_sharding.tile_assignment().num_dimensions() - in PartialReplicateReshardCompatibleSharding()310 (target_sharding.ReplicateOnLastTileDim() ? 1 : 0); in PartialReplicateReshardCompatibleSharding()335 int64 target_tile_size = target_sharding.tile_assignment().dim(dim); in PartialReplicateReshardCompatibleSharding()349 if (target_sharding.ReplicateOnLastTileDim()) { in PartialReplicateReshardCompatibleSharding()351 target_sharding.tile_assignment().dimensions().back(); in PartialReplicateReshardCompatibleSharding()364 if (target_sharding.ReplicateOnLastTileDim()) { in PartialReplicateReshardCompatibleSharding()381 target_sharding.ReplicateOnLastTileDim() in PartialReplicateReshardCompatibleSharding()389 target_sharding.tile_assignment().dimensions()); in PartialReplicateReshardCompatibleSharding()392 target_sharding.tile_assignment().Each( in PartialReplicateReshardCompatibleSharding()[all …]
385 const HloSharding& partial_sharding, const HloSharding& target_sharding);