%2 = "tf.Identity"(%arg0) : (tensor<1x3xf32>) -> tensor<1x3xf32>
  func.return %2 : tensor<1x3xf32>
}

// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<1x3xf32>) -> tensor<*xf32> {
  // CHECK:  "tf.BatchFunction"
  // CHECK-SAME: allowed_batch_sizes = [6]
  // CHECK-SAME: batch_timeout_micros = 100000 : i64
  // CHECK-SAME: batching_queue = ""
  // CHECK-SAME: container = ""
  // CHECK-SAME: enable_large_batch_splitting = false

  func.return %2 : tensor<1x3xf32>
}

// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<1x3xf32>) -> tensor<*xf32> attributes {tf.entry_function = {control_outputs = "", inputs = "input:0", outputs = "batch/BatchFunction:0"}} {
  %0 = "tf.VarHandleOp"() {device = "/device:CPU:0", container = "", shared_name = "variable"} : () -> tensor<!tf_type.resource<tensor<1x3xf32>>>

  // CHECK:  "tf._TfrtGetResource"()
  %cst = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> {device = "/CPU:0"} : () -> tensor<1xi32>
  // CHECK:  "tf.BatchFunction"

namespace {

// Returns the ops that should use node name if shared_name is empty.
const llvm::StringSet<>& GetOpsUsingNodeName() {
  static auto* const ops =
      new llvm::StringSet<>({"VariableV2", "Variable", "BatchFunction"});
  return *ops;
}

// Returns the set of ops that we want to generate shared_names for them if
// empty.
const llvm::StringSet<>& GetSharedNameGenerationCompatibleOps() {
  return GetOpsUsingNodeName();
}

  let assemblyFormat = "$input attr-dict";
}

def BatchFunctionOp : TensorflowMlrt_Op<"batch_function", [Pure]> {
  let summary = "Fallback ExecuteOp specialized for tf.BatchFunction.";

  let description = [{
    This kernel executes a variant tf.BatchFunction kernel that supports having
    the `f` attribute as a bytecode function.

    Example:
      %res = tf_mlrt.batch_function(%input, %captured_input)  {

}

def BatchFunctionOp : FallbackAsync_Op<"batch_function",
    [Pure, CoreRT_TypedAttributeTrait]> {
  let summary = "Fallback ExecuteOp specialized for tf.BatchFunction.";
  let description = [{
    This kernel executes a variant tf.BatchFunction kernel that supports having
    the `f` attribute as a BEF function.

    Example:
      %res = tfrt_fallback_async.batch_function device("/CPU:0") @f(%input, %captured_input)  {

  %v = "tf.AddV2"(%y, %z) {__op_key = 4: i32}: (tensor<i32>, tensor<i32>) -> tensor<i32>
  return %w, %u, %v : tensor<i32>, tensor<i32>, tensor<i32>
}

// -----

// Test lowering BatchFunction op.

func.func @batched_function(%x: tensor<?xi32>) -> tensor<?xi32> {
  return %x : tensor<?xi32>
}

// CHECK-LABEL: func @main

    DefaultValuedOptionalAttr<StrAttr, "\"\"">:$metadata
  );

  let results = (outs
    TF_VariantTensor:$handle
  );
}

def TF_BatchFunctionOp : TF_Op<"BatchFunction", [AttrSizedOperandSegments, DeclareOpInterfaceMethods<SymbolUserOpInterface>, Pure]> {
  let summary = [{
Batches all the inputs tensors to the computation done by the function.
  }];

  let description = [{