// CHECK: %[[VAL_0:.*]] = "tf._FusedConv2D"(%arg2, %arg1, %arg0) <{data_format = "NHWC", dilations = [1, 1, 1, 1], epsilon = 0.000000e+00 : f32, explicit_paddings = [], fused_ops = ["BiasAdd"], num_args = 1 : i64, operandSegmentSizes = array<i32: 1, 1, 1, 0>, padding = "SAME", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true}> {TArgs = [f32]} : (tensor<8x32x32x3xf32>, tensor<1x1x3x128xf32>, tensor<128xf32>) -> tensor<*xf32>

    %0 = "tf.VarHandleOp"() {shared_name = "var_0"} : () -> tensor<!tf_type.resource<tensor<2xi32>>>
    %1 = "tf.BatchFunction"(%0) {
        f = @called_by_batch_func,
        operandSegmentSizes = array<i32: 1, 0>,
        batch_timeout_micros = 1000,
        max_batch_size = 8,
        num_batch_threads = 2
      } : (tensor<!tf_type.resource<tensor<2xi32>>>) -> tensor<*xi32>
    return %1 : tensor<*xi32>

               << ") to be equal to 'n' (" << n << ")";
    }
  }

  Block& block = op.getBody().front();

  auto operandSegmentSizes = op.getProperties().operandSegmentSizes;
  const int32_t num_replicated_inputs = operandSegmentSizes[0];
  const int32_t num_packed_inputs = operandSegmentSizes[1];

  if (num_replicated_inputs % n != 0)
    return op.emitOpError()

    }) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : () -> (tensor<!tf_type.string>, tensor<3x!tf_type.string>)
    tf_device.return %0#0, %0#1: tensor<!tf_type.string>, tensor<3x!tf_type.string>
  }) {n = 2: i32, operandSegmentSizes = array<i32: 0, 0>} : () -> (
      tensor<!tf_type.string>,
      tensor<!tf_type.string>,
      tensor<3x!tf_type.string>,
      tensor<3x!tf_type.string>
      )

      %outputs_2, %control_2 = tf_executor.island wraps "tf.BatchFunction"(%outputs_1) {batch_timeout_micros = 5000 : i64, batching_queue = "", container = "", f = @batched_func, max_batch_size = 256 : i64, num_batch_threads = 2 : i64, operandSegmentSizes = array<i32: 0, 1>, shared_name = ""} : (tensor<*x!tf_type.resource>) -> tensor<*xf32>
      tf_executor.fetch %outputs_2 : tensor<*xf32>
    }
    func.return %0 : tensor<*xf32>
  }

  %2 = "tf.XlaLaunch"(%arg0, %1) {_noinline = true, _xla_compile_device_type = "GPU", device = "/device:GPU:0", function = @xla_func_0, operandSegmentSizes = array<i32: 0, 2, 0>} : (tensor<1x3xf32>, tensor<1x3xf32>) -> tensor<1x3xf32>
  func.return %2 : tensor<1x3xf32>

# RUN: tf-mlir-translate -graphdef-to-mlir -tf-enable-shape-inference-on-import=false -tf-output-arrays=result %s | FileCheck %s

# CHECK: %[[output:.*]], %[[control:.*]]tf_executor.island wraps "tf.ParseExample"
# CHECK: operandSegmentSizes = array<i32: 1, 1, 0, 1, 1>
# CHECK: resultSegmentSizes = array<i32: 0, 0, 0, 1>
# CHECK: tf_executor.fetch %[[output]] : tensor<*xi64>

node {
  name: "serilaized"
  op: "Const"
  attr {
    key: "dtype"

  // CHECK-NEXT: %0 = "tf.ReadVariableOp"(%arg1)
  // CHECK:      [[key:%.*]], [[exec_result:%.*]] = "tf.TPUCompileMlirAndExecute"(%arg0, %0) <{metadata = "metadata", mlir_module = "mlir_module", operandSegmentSizes = array<i32: 2, 0>, operands_with_static_shape = [], producer_name = "default"}> : (tensor<*xi32>, tensor<*xi32>) -> (tensor<3x!tf_type.string>, tensor<*xi32>)
  // CHECK-NEXT: return [[exec_result]] : tensor<*xi32>

  %2 = "tf.XlaLaunch"(%arg0, %1) {_noinline = true, _xla_compile_device_type = "GPU", device = "/device:GPU:0", function = @xla_func_0, operandSegmentSizes = array<i32: 0, 2, 0>} : (tensor<7xf32>, tensor<f32>) -> tensor<10xf32>
  func.return %2 : tensor<10xf32>

    %1:1 = tf_executor.island wraps "tf.SendTPUEmbeddingGradients"(%0) {N = 1 : i64, NN = 0 : i64, config = "test_config_send_embedding", operandSegmentSizes = array<i32: 1, 0>} : (tensor<512x256xf32>) -> () loc("SendTPUEmbedding")
    tf_executor.fetch
  }
  func.return
}

// CHECK:       name: "RecvTPUEmbedding"
// CHECK-NEXT:  op: "RecvTPUEmbeddingActivations"