Searched refs:sharding_tile_dims_stack (Results 1 – 1 of 1) sorted by relevance
318 std::vector<int64> sharding_tile_dims_stack(source_shape.rank()); in ReshapeSharding() local321 sharding_tile_dims_stack[i] = in ReshapeSharding()329 if (Product(sharding_tile_dims_stack) != 1) { in ReshapeSharding()340 s_partitions = sharding_tile_dims_stack.back(); in ReshapeSharding()341 sharding_tile_dims_stack.pop_back(); in ReshapeSharding()345 if (s_partitions * Product(sharding_tile_dims_stack) == 1) { in ReshapeSharding()357 sharding_tile_dims_stack.push_back(s_partitions); in ReshapeSharding()373 sharding_tile_dims_stack.push_back(1); in ReshapeSharding()378 sharding_tile_dims_stack.push_back(s_partitions / t_size); in ReshapeSharding()389 if (sharding_tile_dims_stack.back() != 1 && s_size != s_partitions) { in ReshapeSharding()[all …]