// CHECK-SAME: Tcaptured = [!corert.resource]
      // CHECK-SAME: Tin = []
      // CHECK-SAME: Tout = [f32]

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

  func.func @f(%arg: tensor<1xf32>) -> (tensor<*xf32>, tensor<*xf32>) {
    %handle = "tf.VarHandleOp"() {shared_name = "var1"} : () -> tensor<!tf_type.resource<tensor<0xf32>>>

  // CHECK: [[batch_result:%.*]] = tf_mlrt.await [[batch_result_future]]
  // CHECK-NEXT: [[rendezvous_key_base:%.*]] = tf_mlrt_tpu.compile_and_execute([[batch_result]])
  // CHECK-NEXT: return [[rendezvous_key_base]]

  func.return %arg0 : tensor<f32>
}

// -----

func.func @test_batch_function_with_valid_symbol(%arg0: tensor<1x3xf32>, %arg1: tensor<!tf_type.resource<tensor<1x3xf32>>>) -> () {

    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 = [{