$tensor: A tensorflow tensor.
  }];

  let arguments = (ins
    MlrtPromiseType:$promise,
    TF_Tensor:$tensor
  );
}

def TFMapFnOp : TensorflowMlrt_Op<"tf_map_fn", [AttrSizedOperandSegments, Pure]> {
  let summary = "The Parallel Map for tf_mlrt dialect";
  let description = [{
    The Pmap executes body function in parallel for all ranges up to $max_iterations.

    The pseudo code: