Searched refs:target_tile_assignment_dimensions (Results 1 – 2 of 2) sorted by relevance
749 std::vector<int64> target_tile_assignment_dimensions; in InferShardingFromOperands() local754 target_tile_assignment_dimensions.push_back(1); in InferShardingFromOperands()757 target_tile_assignment_dimensions.push_back( in InferShardingFromOperands()762 target_tile_assignment_dimensions.push_back( in InferShardingFromOperands()766 new_tile_assignment.Reshape(target_tile_assignment_dimensions); in InferShardingFromOperands()1321 std::vector<int64> target_tile_assignment_dimensions( in GetShardingFromUser() local1326 for (int64 i = 0; i < target_tile_assignment_dimensions.size(); ++i) { in GetShardingFromUser()1328 target_tile_assignment_dimensions[i] = in GetShardingFromUser()1331 target_tile_assignment_dimensions[i] = 1; in GetShardingFromUser()1335 tile_assignment.Reshape(target_tile_assignment_dimensions); in GetShardingFromUser()
315 std::vector<int64> target_tile_assignment_dimensions; in ReshapeSharding() local347 target_tile_assignment_dimensions.push_back(1); in ReshapeSharding()352 target_tile_assignment_dimensions.push_back(s_partitions); in ReshapeSharding()355 target_tile_assignment_dimensions.push_back(1); in ReshapeSharding()370 target_tile_assignment_dimensions.push_back(s_partitions); in ReshapeSharding()375 target_tile_assignment_dimensions.push_back(t_size); in ReshapeSharding()402 target_tile_assignment_dimensions.push_back( in ReshapeSharding()405 new_tile_assignment.Reshape(target_tile_assignment_dimensions); in ReshapeSharding()