Searched refs:split_dimension_map (Results 1 – 1 of 1) sorted by relevance
608 const xla::OpSharding& sharding, std::map<int, int>* split_dimension_map) { in GetDimensionIndicesAndNumSplitsFromSharding() argument615 split_dimension_map->emplace( in GetDimensionIndicesAndNumSplitsFromSharding()620 if (split_dimension_map->empty()) { in GetDimensionIndicesAndNumSplitsFromSharding()816 std::map<int, int> split_dimension_map; in CreateOrGetSplitNodesForInputSharding() local818 sharding, &split_dimension_map)); in CreateOrGetSplitNodesForInputSharding()819 TF_RET_CHECK(!split_dimension_map.empty()) in CreateOrGetSplitNodesForInputSharding()835 auto sharding_it = split_dimension_map.begin(); in CreateOrGetSplitNodesForInputSharding()859 while (sharding_it != split_dimension_map.end()) { in CreateOrGetSplitNodesForInputSharding()1237 std::map<int, int> split_dimension_map; in CreateConcatNodesForRetval() local1239 sharding, &split_dimension_map)); in CreateConcatNodesForRetval()[all …]