Searched refs:split_dimension_map (Results 1 – 1 of 1) sorted by relevance
605 const xla::OpSharding& sharding, std::map<int, int>* split_dimension_map) { in GetDimensionIndicesAndNumSplitsFromSharding() argument612 split_dimension_map->emplace( in GetDimensionIndicesAndNumSplitsFromSharding()617 if (split_dimension_map->empty()) { in GetDimensionIndicesAndNumSplitsFromSharding()805 std::map<int, int> split_dimension_map; in CreateOrGetSplitNodesForInputSharding() local807 sharding, &split_dimension_map)); in CreateOrGetSplitNodesForInputSharding()808 TF_RET_CHECK(!split_dimension_map.empty()) in CreateOrGetSplitNodesForInputSharding()824 auto sharding_it = split_dimension_map.begin(); in CreateOrGetSplitNodesForInputSharding()848 while (sharding_it != split_dimension_map.end()) { in CreateOrGetSplitNodesForInputSharding()1024 std::map<int, int> split_dimension_map; in CreateConcatNodesForRetval() local1026 sharding, &split_dimension_map)); in CreateConcatNodesForRetval()[all …]