// CHECK-NEXT:     [[HANDLEY:%.*]] = "tf.VarHandleOp"() <{container = "", shared_name = "y"}> : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>
// CHECK-NEXT:     "tf.IfrtRestoreVariableOp"([[PREFIX]], [[NAME]], [[SLICE]], [[HANDLEY]])
// CHECK-SAME:        {restored_dtypes = [f32], truncate_in_cast = array<i1: false>}
// CHECK-NOT:       "tf.RestoreV2"
// CHECK-NEXT:     return

module {

  );

  let results = (outs
    TFTensorType:$array_key,
    MlrtFutureType: $tensor_future
  );
}

def IfrtRestoreVariableOp: TensorflowMlrt_Op<"ifrt_restore_variable", []> {
  let summary = "Restore variable tensors";
  let description = [{
    This is the MLRT version of tf.IfrtRestoreVariableOp.

    This Op is similar to a combination of RestoreV2 and AssignVariable Op, but
    this Op's execution is asynchronous.

    StrArrayAttr:$shared_name,
    StrArrayAttr:$container
  );

  let results = (outs
    Variadic<TF_Tensor>:$results
  );

  let hasVerifier = 1;
}

def TF_IfrtRestoreVariableOp : TF_Op<"IfrtRestoreVariableOp", []> {
  let summary = "Restore variable tensors";
  let description = [{
    This Op is similar to a combination of RestoreV2 and AssignVariable Op, but
    this Op's execution is asynchronous.

  // CHECK-NEXT: mlrt.await_all_control
  // CHECK-NEXT: return
  func.return
}

// -----

// Test lowering of IfrtRestoreVariableOp

// CHECK-LABEL: func @ifrt_restore_variable_test
func.func @ifrt_restore_variable_test() -> () {
  // CHECK-NEXT: [[PREFIX:%.*]] = tf_mlrt.constop