Searched refs:source_dims_stack (Results 1 – 1 of 1) sorted by relevance
316 std::vector<int64> source_dims_stack(source_shape.rank()); in ReshapeSharding() local320 source_dims_stack[i] = source_shape.dimensions(source_shape.rank() - 1 - i); in ReshapeSharding()327 while (!source_dims_stack.empty() || !target_dims_stack.empty()) { in ReshapeSharding()337 if (!source_dims_stack.empty()) { in ReshapeSharding()338 s_size = source_dims_stack.back(); in ReshapeSharding()339 source_dims_stack.pop_back(); in ReshapeSharding()356 source_dims_stack.push_back(s_size); in ReshapeSharding()372 source_dims_stack.push_back(s_size / t_size); in ReshapeSharding()377 source_dims_stack.push_back(s_size / t_size); in ReshapeSharding()388 CHECK(!source_dims_stack.empty()); in ReshapeSharding()[all …]