// CHECK-SAME:       <{used_by_host = true}> : (tensor<!tf_type.resource<tensor<3x1xf32>>>) -> (tensor<!tf_type.string>, !mlrt.future)
// CHECK-NEXT:    [[TENSOR:%.*]] = "tf_mlrt.tf_await"([[FURTURE]]) : (!mlrt.future) -> tensor<3x1xf32>
// CHECK-NEXT:    "tf.MatMul"(%arg0, [[TENSOR]]) : (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>

include "tfrt/compiler/opdefs/tfrt_op_interfaces.td"  // from @tf_runtime
include "tfrt/compiler/opdefs/tfrt_traits.td"  // from @tf_runtime

// tf_mlrt.tf_await returns a tensorflow Tensor. It is a fake op that is only
// used during parallelization and has no runtime implementation.
def TFAwaitOp: TensorflowMlrt_Op<"tf_await", [Pure, TFRT_CostFunctionInterface, TFRT_FixedCost<1>]> {
  let summary = "Await a tensor from a !mlrt.future";

  let description = [{