Merge pull request tensorflow#52497 from offscale:tensorflow.compiler…

….xla.client.lib.prng.cc PiperOrigin-RevId: 404455831 Change-Id: Icf9fd1e26b80b844e8bbca5c7ee1107cecae8d74
nouiz · Oct 20, 2021 · f65e5d4 · f65e5d4
2 parents 498cb7b + b772327
commit f65e5d4
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/tensorflow/compiler/xla/client/lib/prng.cc b/tensorflow/compiler/xla/client/lib/prng.cc
@@ -195,7 +195,10 @@ SplitShapePair SplitShapeIntoHalves(const Shape& shape) {
   CHECK_GE(pair.split_dim, 0);
   std::vector<int64_t> half_shape_dims;
   std::vector<int64_t> concat_shape_dims;
-  for (int64_t i = 0; i < shape.rank(); ++i) {
+  const auto rank = shape.rank();
+  half_shape_dims.reserve(rank + 1);
+  concat_shape_dims.reserve(rank + 1);
+  for (int64_t i = 0; i < rank; ++i) {
     if (i == pair.split_dim) {
       // Create a new trivial dim for the later concat, which is more friendly
       // to sharding propagation.