diff options
Diffstat (limited to 'tensorflow/core/grappler/optimizers/memory_optimizer.cc')
| -rw-r--r-- | tensorflow/core/grappler/optimizers/memory_optimizer.cc | 11 |
1 files changed, 8 insertions, 3 deletions
diff --git a/tensorflow/core/grappler/optimizers/memory_optimizer.cc b/tensorflow/core/grappler/optimizers/memory_optimizer.cc index 2d1a34f2dd..c775a26914 100644 --- a/tensorflow/core/grappler/optimizers/memory_optimizer.cc +++ b/tensorflow/core/grappler/optimizers/memory_optimizer.cc @@ -37,6 +37,7 @@ limitations under the License. #include "tensorflow/core/grappler/utils.h" #include "tensorflow/core/grappler/utils/topological_sort.h" #include "tensorflow/core/grappler/utils/traversal.h" +#include "tensorflow/core/lib/math/math_util.h" #include "tensorflow/core/protobuf/rewriter_config.pb.h" namespace tensorflow { @@ -1070,9 +1071,13 @@ static bool IdentifySwappingCandidates( // ensure that swapping the tensor back in won't recreate the memory // bottleneck. Last but not least, we want the tensor to have as few // remaining uses as possible. - mem_info.fitness = std::pow((earliest_use - peak_time).count(), 2); - mem_info.fitness /= std::pow(mem_info.uses_left.size(), 2); - mem_info.fitness += std::pow((allocation_time - peak_time).count(), 2); + // + // Note that we must perform the arithmetic inexactly as "double", since + // the values do not fit into any integral type. + mem_info.fitness = + MathUtil::IPow<double>((earliest_use - peak_time).count(), 2) / + MathUtil::IPow<double>(mem_info.uses_left.size(), 2) + + MathUtil::IPow<double>((allocation_time - peak_time).count(), 2); mem_info.fitness = -mem_info.fitness; mem_state.push_back(mem_info); } |