%0 = "tf.TensorListReserve"(%cst_1, %cst_2) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<*xf32>>>
  // CHECK: [[map_fn_result:%.*]] = tf_mlrt.tf_map_fn([[reserve_size]], [[tensor_list]], %arg0)
  // CHECK-SAME: {body_fn = @"map/while_body/MapFnBody", num_tensor_list_or_flow_in = 1 : i32}
  // CHECK-NOT: tf.While

    $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:

  %max_iter = "tf.Const"() {__op_key = 1, value = dense<2> : tensor<i32>} : () -> tensor<i32>
  // CHECK: tf_mlrt.map_fn
  %result = "tf_mlrt.tf_map_fn"(%max_iter, %input2, %2) { operandSegmentSizes = array<i32: 1, 1, 1>, body_fn = @NopMapFnBody, num_tensor_list_or_flow_in = 1 : i32} : (tensor<i32>, tensor<!tf_type.variant<tensor<*xf32>>>, tensor<i32>) -> tensor<i32>