1 files changed, 4 insertions, 8 deletions
diff --git a/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.cc b/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.cc
index 114a9241bd..00ccfb1c78 100644
--- a/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.cc
+++ b/tensorflow/compiler/tf2xla/xla_jit_compiled_cpu_function.cc
@@ -58,15 +58,11 @@ xla::StatusOr<std::vector<intptr_t>> ComputeTempSizes(
   std::vector<intptr_t> temp_sizes;
   temp_sizes.reserve(allocations.size());
   for (const xla::BufferAllocation& allocation : allocations) {
-    if (allocation.is_constant() || allocation.is_thread_local()) {
-      // Constants are lowered to globals.  Thread locals are lowered to
-      // allocas.
+    // Callers don't allocate temporary buffers for parameters. Nor for
+    // thread-local buffers, which are lowered to alloca.
+    if (allocation.is_entry_computation_parameter() ||
+        allocation.is_thread_local()) {
       temp_sizes.push_back(-1);
-    } else if (allocation.is_entry_computation_parameter()) {
-      // Entry computation parameters need some preprocessing in
-      // XlaCompiledCpuFunction::Run.  See the comment on
-      // XlaCompiledCpuFunction::StaticData::temp_sizes.
-      temp_sizes.push_back(-allocation.parameter_number() - 2);
     } else {
       temp_sizes.push_back(allocation.size());
     }