[pull] master from tensorflow:master

tensorflow/compiler/mlir/lite/tests/optimize_batch_matmul.mlir

@@ -342,4 +342,68 @@ func.func @Batchmatmul2FullyconnectedReshapeTransposeChainConstY(%arg0: tensor<1
   // CHECK-NOT: "tfl.batch_matmul"
   %2 = "tfl.batch_matmul"(%arg0, %1) {adj_x = false, adj_y = false, asymmetric_quantize_inputs = false} : (tensor<16x1024xf32>, tensor<1024x128xf32>) -> tensor<16x128xf32>
   func.return %2 : tensor<16x128xf32>
+}
+
+// A bf16 constant widened to f32 is still foldable, since flatbuffer export
+// materializes the cast. The transpose goes below the cast so it lands
+// directly on the constant, where it can be folded.
+// CHECK-LABEL: Batchmatmul2FullyconnectedBf16CastConstY
+func.func @Batchmatmul2FullyconnectedBf16CastConstY(%arg0: tensor<16x1024xf32>) -> tensor<16x128xf32> {
+  %weight = arith.constant dense_resource<__elided__> : tensor<1024x128xbf16>
+  %0 = "tfl.cast"(%weight) : (tensor<1024x128xbf16>) -> tensor<1024x128xf32>
+  // CHECK: %[[W:.*]] = arith.constant dense_resource<__elided__> : tensor<1024x128xbf16>
+  // CHECK: %[[T:.*]] = "tfl.transpose"(%[[W]], %{{.*}}) : (tensor<1024x128xbf16>, tensor<2xi32>) -> tensor<128x1024xbf16>
+  // CHECK: %[[C:.*]] = "tfl.cast"(%[[T]]) : (tensor<128x1024xbf16>) -> tensor<128x1024xf32>
+  // CHECK: "tfl.fully_connected"(%arg0, %[[C]], %{{.*}})
+  // CHECK-NOT: "tfl.batch_matmul"
+  %1 = "tfl.batch_matmul"(%arg0, %0) {adj_x = false, adj_y = false, asymmetric_quantize_inputs = false} : (tensor<16x1024xf32>, tensor<1024x128xf32>) -> tensor<16x128xf32>
+  func.return %1 : tensor<16x128xf32>
+}
+
+// With adj_y the rhs is already in filter layout, so the cast is kept as is.
+// CHECK-LABEL: Batchmatmul2FullyconnectedBf16CastConstYAdjY
+func.func @Batchmatmul2FullyconnectedBf16CastConstYAdjY(%arg0: tensor<16x1024xf32>) -> tensor<16x128xf32> {
+  %weight = arith.constant dense_resource<__elided__> : tensor<128x1024xbf16>
+  %0 = "tfl.cast"(%weight) : (tensor<128x1024xbf16>) -> tensor<128x1024xf32>
+  // CHECK: %[[W:.*]] = arith.constant dense_resource<__elided__> : tensor<128x1024xbf16>
+  // CHECK: %[[C:.*]] = "tfl.cast"(%[[W]]) : (tensor<128x1024xbf16>) -> tensor<128x1024xf32>
+  // CHECK-NOT: "tfl.transpose"
+  // CHECK: "tfl.fully_connected"(%arg0, %[[C]], %{{.*}})
+  // CHECK-NOT: "tfl.batch_matmul"
+  %1 = "tfl.batch_matmul"(%arg0, %0) {adj_x = false, adj_y = true, asymmetric_quantize_inputs = false} : (tensor<16x1024xf32>, tensor<128x1024xf32>) -> tensor<16x128xf32>
+  func.return %1 : tensor<16x128xf32>
+}
+
+// CHECK-LABEL: NotFuseTransposeFCRhsToBatchMatmulBf16CastConstant
+func.func @NotFuseTransposeFCRhsToBatchMatmulBf16CastConstant(%arg0: tensor<16x1024xf32>, %arg1: none) -> tensor<16x128xf32> {
+  %cst = arith.constant dense<[1, 0]> : tensor<2xi32>
+  %weight = arith.constant dense_resource<__elided__> : tensor<1024x128xbf16>
+  %0 = "tfl.cast"(%weight) : (tensor<1024x128xbf16>) -> tensor<1024x128xf32>
+  %1 = "tfl.transpose"(%0, %cst) : (tensor<1024x128xf32>, tensor<2xi32>) -> tensor<128x1024xf32>
+  // CHECK: "tfl.fully_connected"
+  // CHECK-NOT: "tfl.batch_matmul"
+  %2 = "tfl.fully_connected"(%arg0, %1, %arg1) {asymmetric_quantize_inputs = false, fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<16x1024xf32>, tensor<128x1024xf32>, none) -> tensor<16x128xf32>
+  func.return %2 : tensor<16x128xf32>
+}
+
+// A widening cast of a non-constant is not foldable.
+// CHECK-LABEL: NotBatchmatmul2FullyconnectedBf16CastNonConstY
+func.func @NotBatchmatmul2FullyconnectedBf16CastNonConstY(%arg0: tensor<16x1024xf32>, %arg1: tensor<1024x128xbf16>) -> tensor<16x128xf32> {
+  %0 = "tfl.cast"(%arg1) : (tensor<1024x128xbf16>) -> tensor<1024x128xf32>
+  // CHECK: "tfl.batch_matmul"
+  // CHECK-NOT: "tfl.fully_connected"
+  %1 = "tfl.batch_matmul"(%arg0, %0) {adj_x = false, adj_y = false, asymmetric_quantize_inputs = false} : (tensor<16x1024xf32>, tensor<1024x128xf32>) -> tensor<16x128xf32>
+  func.return %1 : tensor<16x128xf32>
+}
+
+// Only float widening is looked through; flatbuffer export does not
+// materialize other casts.
+// CHECK-LABEL: NotBatchmatmul2FullyconnectedI8CastConstY
+func.func @NotBatchmatmul2FullyconnectedI8CastConstY(%arg0: tensor<16x1024xf32>) -> tensor<16x128xf32> {
+  %weight = arith.constant dense_resource<__elided__> : tensor<1024x128xi8>
+  %0 = "tfl.cast"(%weight) : (tensor<1024x128xi8>) -> tensor<1024x128xf32>
+  // CHECK: "tfl.batch_matmul"
+  // CHECK-NOT: "tfl.fully_connected"
+  %1 = "tfl.batch_matmul"(%arg0, %0) {adj_x = false, adj_y = false, asymmetric_quantize_inputs = false} : (tensor<16x1024xf32>, tensor<1024x128xf32>) -> tensor<16x128xf32>
+  func.return %1 : tensor<16x128xf32>
 }

tensorflow/compiler/mlir/lite/transforms/optimize_batch_matmul_pass.cc

@@ -44,9 +44,26 @@ namespace mlir {
 namespace TFL {
 namespace {
 
+// Returns `op` as a bf16/f16 -> f32 cast, or null if it is anything else.
+//
+// Weights kept in a narrow float type reach the matmul behind such a cast,
+// since the kernels only run in f32. Flatbuffer export materializes this cast
+// when its input is a constant, so the cast does not stop the value from being
+// a constant filter.
+CastOp AsFloatWideningCast(mlir::Operation* op) {
+  auto cast_op = mlir::dyn_cast_or_null<CastOp>(op);
+  if (!cast_op) return nullptr;
+  mlir::Type in_elem = getElementTypeOrSelf(cast_op.getInput().getType());
+  mlir::Type out_elem = getElementTypeOrSelf(cast_op.getType());
+  if (!(in_elem.isBF16() || in_elem.isF16()) || !out_elem.isF32()) {
+    return nullptr;
+  }
+  return cast_op;
+}
+
 // Checks whether `value` is, or can be folded into, a constant. This covers a
 // plain constant, the dequantize wrappers a quantized constant weight hides
-// behind, and chains of ops that only rearrange one of those.
+// behind, and chains of ops that only rearrange or widen one of those.
 bool IsFromFoldableChain(mlir::Value value) {
   mlir::Operation* defining_op = value.getDefiningOp();
 
@@ -66,6 +83,8 @@ bool IsFromFoldableChain(mlir::Value value) {
       defining_op = split_op.getValue().getDefiningOp();
     } else if (auto transpose_op = mlir::dyn_cast<TransposeOp>(defining_op)) {
       defining_op = transpose_op.getInput().getDefiningOp();
+    } else if (auto cast_op = AsFloatWideningCast(defining_op)) {
+      defining_op = cast_op.getInput().getDefiningOp();
     } else {
       // Stop if the op does not preserve constant nature.
       break;
@@ -193,8 +212,24 @@ struct ConvertBatchMatMulOp2FullyConnectedOp_Rank2ConstantRhs
 
     // The rhs need to be transposed if adj_y == false AND this matmul will be
     // legalized to tfl.fully_connected
-    Value output_rhs =
-        !bmm_op.getAdjY() ? create_z_x_transpose_op(input_rhs) : input_rhs;
+    Value output_rhs = input_rhs;
+    if (!bmm_op.getAdjY()) {
+      // Transpose below a widening cast rather than above it. Only a
+      // transpose directly on a constant gets folded away, and the cast is
+      // then left directly on the constant too, which flatbuffer export
+      // materializes.
+      if (CastOp cast_op = AsFloatWideningCast(input_rhs.getDefiningOp())) {
+        Value transposed = create_z_x_transpose_op(cast_op.getInput());
+        auto transposed_type =
+            mlir::cast<RankedTensorType>(transposed.getType());
+        output_rhs = TFL::CastOp::create(
+            rewriter, cast_op.getLoc(),
+            transposed_type.clone(getElementTypeOrSelf(cast_op.getType())),
+            transposed);
+      } else {
+        output_rhs = create_z_x_transpose_op(input_rhs);
+      }
+    }
 
     Type output_type = bmm_op.getResult().getType();
     auto no_input =

tensorflow/compiler/mlir/tfr/ir/tfr_ops.td

@@ -41,6 +41,7 @@ def TFR_Dialect : Dialect {
   }];
 
   let cppNamespace = "::mlir::TFR";
+  let useStrictPropertiesInAssemblyFormat = 0;
 }
 
 //===----------------------------------------------------------------------===//

tensorflow/compiler/mlir/tfrt/ir/mlrt/mlrt_dialect.td

@@ -26,6 +26,7 @@ def Mlrt_Dialect : Dialect {
   }];
 
   let cppNamespace = "::mlrt::compiler";
+  let useStrictPropertiesInAssemblyFormat = 0;
 }
 
 def MlrtFutureType : DialectType<Mlrt_Dialect,

tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_dialect.td

@@ -29,6 +29,7 @@ def TensorflowMlrt_Dialect : Dialect {
   }];
 
   let cppNamespace = "::tensorflow::tf_mlrt";
+  let useStrictPropertiesInAssemblyFormat = 0;
 }
 
 class TensorflowMlrt_Op<string mnemonic, list<Trait> traits = []> :

tensorflow/compiler/mlir/tfrt/ir/mlrt/tf_mlrt_tpu_ops.td

@@ -28,6 +28,7 @@ def TensorflowMlrtTpu_Dialect : Dialect {
   }];
 
   let cppNamespace = "::tensorflow::tf_mlrt_tpu";
+  let useStrictPropertiesInAssemblyFormat = 0;
 }
 
 class TensorflowMlrtTpu_Op<string mnemonic, list<Trait> traits = []> :

tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_async.td

@@ -33,6 +33,7 @@ def FallbackAsync_Dialect : Dialect {
   }];
 
   let cppNamespace = "::tfrt::fallback_async";
+  let useStrictPropertiesInAssemblyFormat = 0;
 }
 
 class FallbackAsync_Op<string mnemonic, list<Trait> traits = []> :

tensorflow/compiler/mlir/tfrt/ir/tfrt_fallback_sync.td

@@ -32,6 +32,7 @@ def FallbackSync_Dialect : Dialect {
   }];
 
   let cppNamespace = "::tfrt::fallback_sync";
+  let useStrictPropertiesInAssemblyFormat = 0;
 }
 
 class FallbackSync_Op<string mnemonic, list<Trait> traits = []> :

tensorflow/compiler/mlir/tools/kernel_gen/ir/tf_framework_ops.td

@@ -40,6 +40,7 @@ def TFFramework_Dialect : Dialect {
   }];
 
   let useDefaultTypePrinterParser = 1;
+  let useStrictPropertiesInAssemblyFormat = 1;
 }
 
 def TFFramework_OpKernelContextType : DialectType<TFFramework_Dialect,
@@ -113,7 +114,7 @@ def TFFramework_TFAllocOp : TFFramework_Op<"alloc", [
         "reuse_input_candidates";
   }];
   let assemblyFormat = [{
-    `(` $ctx (`,` $dyn_sizes^ )? `)` attr-dict `:` type($result)
+    `(` $ctx (`,` $dyn_sizes^ )? `)` prop-dict attr-dict `:` type($result)
   }];
 
   let hasVerifier = 1;
@@ -294,7 +295,7 @@ def TFFramework_JITCompileFromStrOp : TFFramework_Op<"jit_compile_from_str",
     static constexpr StringRef kJITEntryFunctionName = "main";
   }];
 
-  let assemblyFormat = "($ctx^ `,`)? $code attr-dict";
+  let assemblyFormat = "($ctx^ `,`)? $code prop-dict attr-dict";
 }
 
 //===----------------------------------------------------------------------===//

tensorflow/compiler/mlir/tools/kernel_gen/tests/embed_tf_framework.mlir

@@ -92,9 +92,9 @@ func.func @jit_compile_from_str(%ctx : !tf_framework.op_kernel_context)
     -> !tf_framework.jit_callable attributes {tf_entry} {
   // CHECK: %[[RES:.*]] = tf_framework.jit_compile_from_str %[[CTX]], "placeholder"
   // CHECK: return %[[RES]]
-  %0 = tf_framework.jit_compile_from_str "placeholder" {
-      architectures = ["sm_123", "sm_456"], tileSizes = [1, 2, 3],
-      unrollFactors = [4], enableFtz = false,
-      index64Bit = false, cpuCodegen = false }
+  %0 = tf_framework.jit_compile_from_str "placeholder" <
+      tileSizes = [1, 2, 3], unrollFactors = [4], enableFtz = false,
+      index64Bit = false, cpuCodegen = false> {
+      architectures = ["sm_123", "sm_456"] }
   func.return %0 : !tf_framework.jit_callable
 }

tensorflow/compiler/mlir/tools/kernel_gen/tests/func_to_jit_invocations.mlir

@@ -41,12 +41,12 @@ func.func @unary_tanh(%arg : tensor<?xf32>) -> tensor<?xf32> {
 // CHECK-SAME:        }
 // CHECK-SAME:      }
 // CHECK-SAME:      "
-// CHECK-SAME:      {
-// CHECK-SAME:        cpuCodegen = false
-// CHECK-SAME:        enableFtz = false
+// CHECK-SAME:      <
 // CHECK-SAME:        tileSizes = [1, 2, 3]
 // CHECK-SAME:        unrollFactors = [3, 2, 1]
-// CHECK-SAME:      }
+// CHECK-SAME:        enableFtz = false
+// CHECK-SAME:        cpuCodegen = false
+// CHECK-SAME:      >
 // CHECK:       %[[RES:.*]] = tf_framework.jit_execute %[[CALLABLE]](%[[ARG]])
 // CHECK:       return %[[RES]]
 
@@ -59,11 +59,11 @@ func.func @unary_tanh(%arg : tensor<?xf32>) -> tensor<?xf32> {
 // CHECK-JFLT:       %[[IF:.*]] = scf.if %[[CMPI]]
 // CHECK-JFLT:         %[[JIT:.*]] = tf_framework.jit_compile_from_str
 // CHECK-JFLT-SAME:        "module
-// CHECK-JFLT-SAME:        cpuCodegen = false
-// CHECK-JFLT-SAME:        enableFtz = false
-// CHECK-JFLT-SAME:        index64Bit = true
 // CHECK-JFLT-SAME:        tileSizes = [1, 2, 3]
 // CHECK-JFLT-SAME:        unrollFactors = [3, 2, 1]
+// CHECK-JFLT-SAME:        enableFtz = false
+// CHECK-JFLT-SAME:        index64Bit = true
+// CHECK-JFLT-SAME:        cpuCodegen = false
 // CHECK-JFLT:         %[[JIT_0:.*]] = tf_framework.jit_execute %[[JIT]](%[[ARG0]])
 // CHECK-JFLT:         scf.yield %[[JIT_0]]
 // CHECK-JFLT:       else
@@ -91,12 +91,12 @@ func.func @binary_sub(%arg0 : tensor<*xf32>, %arg1 : tensor<*xf32>) -> tensor<*x
 // CHECK-SAME:        }
 // CHECK-SAME:      }
 // CHECK-SAME:      "
-// CHECK-SAME:      {
-// CHECK-SAME:        cpuCodegen = false
-// CHECK-SAME:        enableFtz = false
+// CHECK-SAME:      <
 // CHECK-SAME:        tileSizes = [1, 2, 3]
 // CHECK-SAME:        unrollFactors = [3, 2, 1]
-// CHECK-SAME:      }
+// CHECK-SAME:        enableFtz = false
+// CHECK-SAME:        cpuCodegen = false
+// CHECK-SAME:      >
 // CHECK:       %[[RES:.*]] = tf_framework.jit_execute %[[CALLABLE]](%[[ARG0]], %[[ARG1]])
 // CHECK:       return %[[RES]]
 
@@ -122,12 +122,12 @@ func.func @binary_sub(%arg0 : tensor<*xf32>, %arg1 : tensor<*xf32>) -> tensor<*x
 // CHECK-JFLT-SAME:        }
 // CHECK-JFLT-SAME:      }
 // CHECK-JFLT-SAME:      "
-// CHECK-JFLT-SAME:      {
-// CHECK-JFLT-SAME:        cpuCodegen = false
-// CHECK-JFLT-SAME:        enableFtz = false
+// CHECK-JFLT-SAME:      <
 // CHECK-JFLT-SAME:        tileSizes = [1, 2, 3]
 // CHECK-JFLT-SAME:        unrollFactors = [3, 2, 1]
-// CHECK-JFLT-SAME:      }
+// CHECK-JFLT-SAME:        enableFtz = false
+// CHECK-JFLT-SAME:        cpuCodegen = false
+// CHECK-JFLT-SAME:      >
 // CHECK-JFLT:       %[[RES:.*]] = tf_framework.jit_execute %[[CALLABLE]](%[[ARG0]], %[[ARG1]])
 // CHECK-JFLT:       scf.yield %[[RES]] : tensor<*xf32>
 // CHECK-JFLT:     } else {
@@ -155,13 +155,13 @@ func.func @reciprocal(%arg0: tensor<*xf32>)
 // CHECK=SAME:       return %[[RES_JIT]] : tensor<*xf32>
 // CHECK-SAME:     }
 // CHECK-SAME:   }
-// CHECK-SAME:   {
-// CHECK-SAME:     cpuCodegen = false,
+// CHECK-SAME:   <
+// CHECK-SAME:     tileSizes = [1, 2, 3],
+// CHECK-SAME:     unrollFactors = [3, 2, 1],
 // CHECK-SAME:     enableFtz = false,
 // CHECK-SAME:     index64Bit = false,
-// CHECK-SAME:     tileSizes = [1, 2, 3],
-// CHECK-SAME:     unrollFactors = [3, 2, 1]
-// CHECK-SAME:   }
+// CHECK-SAME:     cpuCodegen = false
+// CHECK-SAME:   >
 // CHECK:      %[[RES:.*]] = tf_framework.jit_execute %[[CALLABLE]](%[[ARG]]) : tensor<*xf32> -> tensor<*xf32>
 // CHECK:      return %[[RES]] : tensor<*xf32>
 
@@ -174,11 +174,11 @@ func.func @reciprocal(%arg0: tensor<*xf32>)
 // CHECK-JFLT:       %[[IF:.*]] = scf.if %[[CMPI]]
 // CHECK-JFLT:         %[[JIT:.*]] = tf_framework.jit_compile_from_str
 // CHECK-JFLT-SAME:        "module
-// CHECK-JFLT-SAME:        cpuCodegen = false
-// CHECK-JFLT-SAME:        enableFtz = false
-// CHECK-JFLT-SAME:        index64Bit = true
 // CHECK-JFLT-SAME:        tileSizes = [1, 2, 3]
 // CHECK-JFLT-SAME:        unrollFactors = [3, 2, 1]
+// CHECK-JFLT-SAME:        enableFtz = false
+// CHECK-JFLT-SAME:        index64Bit = true
+// CHECK-JFLT-SAME:        cpuCodegen = false
 // CHECK-JFLT:         %[[JIT_0:.*]] = tf_framework.jit_execute %[[JIT]](%[[ARG0]])
 // CHECK-JFLT:         scf.yield %[[JIT_0]]
 // CHECK-JFLT:       else

tensorflow/compiler/mlir/tools/kernel_gen/tests/ops.mlir

@@ -29,10 +29,10 @@ func.func @alloc(%ctx: !tf_framework.op_kernel_context,
 // CHECK-LABEL: func @forwarding_alloc
 func.func @forwarding_alloc(%ctx: !tf_framework.op_kernel_context,
                        %size_0 : index , %size_2 : index) {
-  %buf = tf_framework.alloc(%ctx, %size_0, %size_2) {
+  %buf = tf_framework.alloc(%ctx, %size_0, %size_2) <
     input_indices = [0 : i32, 1 : i32],
-    output_index = 0 : i32
-  } : memref<?x10x?xi8>
+    output_index = 0
+  > : memref<?x10x?xi8>
   func.return
 }
 
@@ -94,20 +94,20 @@ func.func @jit_compile(%ctx : !tf_framework.op_kernel_context)
 
 // CHECK-LABEL: func @jit_compile_from_str_wo_ctx
 func.func @jit_compile_from_str_wo_ctx() -> !tf_framework.jit_callable {
-  %callable = tf_framework.jit_compile_from_str "placeholder" {
-      architectures = ["sm_123", "sm_456"], tileSizes = [1, 2, 3],
-      unrollFactors = [4], enableFtz = false,
-      index64Bit = false, cpuCodegen = false }
+  %callable = tf_framework.jit_compile_from_str "placeholder" <
+      tileSizes = [1, 2, 3], unrollFactors = [4], enableFtz = false,
+      index64Bit = false, cpuCodegen = false> {
+      architectures = ["sm_123", "sm_456"] }
   func.return %callable : !tf_framework.jit_callable
 }
 
 // CHECK-LABEL: func @jit_compile_from_str
 func.func @jit_compile_from_str(%ctx : !tf_framework.op_kernel_context)
     -> !tf_framework.jit_callable {
-  %callable = tf_framework.jit_compile_from_str %ctx , "placeholder" {
-      architectures = ["sm_123", "sm_456"], tileSizes = [1, 2, 3],
-      unrollFactors = [4], enableFtz = false,
-      index64Bit = false, cpuCodegen = false }
+  %callable = tf_framework.jit_compile_from_str %ctx , "placeholder" <
+      tileSizes = [1, 2, 3], unrollFactors = [4], enableFtz = false,
+      index64Bit = false, cpuCodegen = false> {
+      architectures = ["sm_123", "sm_456"] }
   func.return %callable : !tf_framework.jit_callable
 }
 

tensorflow/compiler/mlir/tools/kernel_gen/tests/tf_framework_legalize_to_llvm.mlir

@@ -206,9 +206,9 @@ func.func @jit_compile_from_str(%ctx: !tf_framework.op_kernel_context)
   // CHECK-SAME: %[[NUM_UNROLL_FACTORS]], %[[UNROLL_FACTORS]],
   // CHECK-SAME: %[[ENABLE_FTZ]], %[[CPU_CODEGEN]]
   // CHECK: llvm.return %[[RES]]
-  %0 = tf_framework.jit_compile_from_str %ctx, "placeholder" {
+  %0 = tf_framework.jit_compile_from_str %ctx, "placeholder" <
       tileSizes = [1, 2, 3], unrollFactors = [4],
-      enableFtz = false, index64Bit = false, cpuCodegen = false }
+      enableFtz = false, index64Bit = false, cpuCodegen = false>
   func.return %0 : !tf_framework.jit_callable
 }
 

tensorflow/core/common_runtime/rendezvous_util.cc

@@ -75,14 +75,14 @@ void RecvOutputsFromRendezvousAsync(
         keys.size(), "; alloc_attrs.size() = ", alloc_attrs.size())));
   }
 
-  received_tensors->reserve(keys.size());
+  received_tensors->resize(keys.size());
   std::vector<std::tuple<std::string, Tensor*, Rendezvous::ParsedKey,
                          AllocatorAttributes>>
       arguments;
+  arguments.reserve(keys.size());
   for (int i = 0; i < keys.size(); ++i) {
     Rendezvous::ParsedKey parsed;
     absl::Status s = Rendezvous::ParseKey(keys[i], &parsed);
-    received_tensors->push_back(Tensor());
     if (!s.ok()) {
       done(s);
       return;