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 = [{
    Await a tensor from a !mlrt.future.

def TensorflowMlrt_Dialect : Dialect {
  let name = "tf_mlrt";

  let description = [{
    The TF MLRT Dialect.
  }];

  let cppNamespace = "::tensorflow::tf_mlrt";
}

class TensorflowMlrt_Op<string mnemonic, list<Trait> traits = []> :
    Op<TensorflowMlrt_Dialect, mnemonic, !listconcat(traits, [IsolatedFromAbove])> {
}

// This corresponds to tensorflow::Tensor.