func.return %0, %1 : tensor<*xf32>, tensor<*xf32>
  }
  // CHECK: func.func @f(%[[ARG_1:.*]]: tensor<1xf32>)
  // Make sure that `operandSegmentSizes` attribute is also updated.

        finalize_func = @__finalize_func_790,
        init_func = @__init_func_530, next_func = @__next_func_680,
        next_func.experimental_ints_on_device = true,
        operandSegmentSizes = array<i32: 2, 2, 1>,
        output_shapes = [#tf_type.shape<?>],
        output_types = [f32],
        metadata = ""
      } :

      %0 = "tf.RecvTPUEmbeddingActivations"() {config = "test_config_recv_embedding"} : () -> tensor<512x256xf32>
      "tf.SendTPUEmbeddingGradients"(%0) {N = 1 : i64, NN = 0 : i64, config = "test_config_send_embedding", operandSegmentSizes = array<i32: 1, 0>} : (tensor<512x256xf32>) -> ()
      "tf.A"() {_xla_outside_compilation = "cluster1"} : () -> ()
      tf_device.return

      device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0",
      finalize_func = @finalize,
      init_func = @init,
      next_func = @next,
      operandSegmentSizes = array<i32: 1, 2, 1>,
      output_shapes = [#tf_type.shape<>],
      output_types = [!tf_type.string],
      metadata = ""} : (
              tensor<4xf32>,
              tensor<3xf32>,

%0, %0, %0, %0, %0, %0) {dense_shapes = [#tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<>, #tf_type.shape<>], device = "/job:localhost/replica:0/task:0/device:CPU:0", operandSegmentSizes = array<i32: 1, 1, 0, 8, 8>, resultSegmentSizes = array<i32: 0, 0, 0, 8>} : (tensor<?x!tf_type.string>, tensor<0x!tf_type.string>, tensor<!tf_type.string>, tensor<!tf_type.string>, tensor<!tf_type.string>, tensor<!tf_type.string>, tensor<!tf_type.string>,...

        finalize_func = @__finalize_func_790,
        init_func = @__init_func_530, next_func = @__next_func_680,
        next_func.experimental_ints_on_device = true,
        operandSegmentSizes = array<i32: 2, 2, 1>,
        output_shapes = [#tf_type.shape<?>],
        output_types = [f32],
        metadata = ""
      } :

    data_format = "NHWC", dilations = [1, 1, 1, 1], epsilon = 9.99999974E-5 : f32, explicit_paddings = [], filter_format = "HWIO", fused_ops = ["BiasAdd", "Relu"], leakyrelu_alpha = 2.000000e-01 : f32, num_args = 2 : i64, operandSegmentSizes = array<i32: 1, 1, 2, 2>, padding = "SAME", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true
    } : (tensor<1x300x300x40xi8>, tensor<3x3x40x40xi8>, tensor<40xf32>, tensor<0xi8>, tensor<f32>, tensor<f32>) -> tensor<1x300x300x40xi8>

    %1 = func.call @finalize(%arg4, %arg5) : (tensor<6xf32>, tensor<2xf32>) -> tensor<5xf32>
    "tf.Yield"(%1) : (tensor<5xf32>) -> ()
  }) {device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0", metadata = "", operandSegmentSizes = array<i32: 1, 2, 1>, output_shapes = [#tf_type.shape<>], output_types = [!tf_type.string]} : (tensor<4xf32>, tensor<3xf32>, tensor<!tf_type.resource>, tensor<2xf32>) -> tensor<!tf_type.variant>
  return
}

// -----

    "tf.SendTPUEmbeddingGradients"(%2#0, %2#1) {_tpu_embedding_layer = "call1", config = "\0A\0B\0C\0D", operandSegmentSizes = array<i32: 2, 0>} : (tensor<2x2xf32>, tensor<4x4xf32>) -> ()
    tf_device.return
  }) {allow_soft_placement = true, num_cores_per_replica = 1, topology =  "", device_assignment =  []} : () -> ()
  func.return
}

}

void BatchFunctionOp::eraseArguments(const BitVector& erase_indices) {
  const StringRef operand_segment_size_attr = getOperandSegmentSizeAttr();
  auto operandSegmentSizes = getOperation()->getAttrOfType<DenseI32ArrayAttr>(
      operand_segment_size_attr);

  // `operandSegmentSizes` attribute indicates the sizes of the two
  // variadic operands of `BatchFunctionOp`: `in_tensors` and