Searched refs:source_sharding (Results 1 – 3 of 3) sorted by relevance
1796 const HloSharding& source_sharding, in CreateMatchingShardingOnDims() argument1801 if (source_sharding.IsReplicated()) { in CreateMatchingShardingOnDims()1807 num_tiles *= source_sharding.tile_assignment().dim(source_dims[i]); in CreateMatchingShardingOnDims()1809 source_sharding.tile_assignment().dim(source_dims[i]); in CreateMatchingShardingOnDims()1814 if (num_tiles != source_sharding.tile_assignment().num_elements()) { in CreateMatchingShardingOnDims()1815 CHECK_EQ(source_sharding.tile_assignment().num_elements() % num_tiles, 0); in CreateMatchingShardingOnDims()1817 tile_dims.push_back(source_sharding.tile_assignment().num_elements() / in CreateMatchingShardingOnDims()1820 auto tgt_tile_assignment = source_sharding.tile_assignment(); in CreateMatchingShardingOnDims()1824 target_dims, source_sharding, source_dims); in CreateMatchingShardingOnDims()1827 target_dims, source_sharding, source_dims); in CreateMatchingShardingOnDims()
409 const HloSharding& source_sharding,
2480 const auto source_sharding = in HandleGetTupleElement() local2482 gte->set_sharding(source_sharding); in HandleGetTupleElement()