return
}

// -----

// CHECK-LABEL: @get_value
func.func @get_value(%arg0: tensor<!tf_type.variant<tensor<f32>>>) {
  // CHECK-NOT: Optional
  %1 = "tf.OptionalGetValue"(%arg0) : (tensor<!tf_type.variant<tensor<f32>>>) -> tensor<f32>
  return
}

// -----

// CHECK-LABEL: @none
func.func @none(%arg0: tensor<!tf_type.variant<tensor<f32>>>) {
  // CHECK-NOT: Optional

    // CHECK-NEXT:     %[[HAS_VALUE:[0-9]*]] = "tf.OptionalHasValue"(%[[GET_NEXT]])
    // CHECK-NEXT:     %[[IF:.*]] = "tf.IfRegion"(%[[HAS_VALUE]])
    // CHECK-NEXT:       %[[GET_VALUE:[0-9]*]] = "tf.OptionalGetValue"(%[[GET_NEXT]])
    // CHECK-NEXT:       %[[FUNC_CALL:[0-9]*]] = func.call @__reduce_func_1(%[[ARG_5]], %[[GET_VALUE]])
    // CHECK-SAME:       _xla_compile_device_type = "TPU"

  // their "no-op" expansion is known to cause problems in some cases.
  static const char* kOpsToSkip[] = {
      "IdentityOp",
      "NoOp",              // b/174596063
      "OptionalHasValue",  // b/173136483
      "OptionalGetValue",  // b/173136483
      "VarHandleOp",       // b/176819198
  };
  for (const char* skip : kOpsToSkip) {
    if (absl::StartsWith(orig_op->op_name(), skip)) return absl::OkStatus();
  }

    auto input = op.getOptional();
    auto elementType = getElementTypeOrSelf(input.getType());
    if (isa<TF::VariantType>(elementType)) {
      // We can only replace OptionalGetValue after the inputs have been
      // replaced.
      return failure();
    }
    rewriter.replaceOpWithNewOp<TF::CastOp>(op, op.getResult(0).getType(),
                                            input);

      Name("OptionalHasValue").Device(DEVICE).HostMemory("has_value"),         \
      data::OptionalHasValueOp);                                               \
  REGISTER_KERNEL_BUILDER(Name("OptionalGetValue").Device(DEVICE),             \
                          data::OptionalGetValueOp);                           \
  REGISTER_KERNEL_BUILDER(Name(FunctionLibraryDefinition::kArgOp)              \

    raw_ostream* os = nullptr);

// Moves TPUCompileMlir ops as far to the front as possible.
std::unique_ptr<OperationPass<func::FuncOp>> CreateMoveTpuCompileToFrontPass();

// Decomposes OptionalFromValue, OptionalGetValue, OptionalNone,
// and OptionalHasValue
std::unique_ptr<OperationPass<ModuleOp>> CreateDecomposeOptionalsPass();

//===----------------------------------------------------------------------===//
// XlaCallModule

  let summary = "Decomposes ops of the 'Optional' family";
  let description = [{
    Turns OptionalFromValue, OptionalGetValue, OptionalNone,
    and OptionalHasValue into identities.

    For example, if we have:

      %1 = "tf.OptionalFromValue"(%0) : (tensor<f32>) -> !tf_type.variant
      %2 = "tf.OptionalGetValue"(%1) : (tensor<!tf_type.variant>) -> tensor<f32>

    then after the pass, this will have been transformed to:

    %3 = "tf.OptionalHasValue"(%2) : (tensor<!tf_type.variant>) -> tensor<i1>
    %4 = "tf.IfRegion"(%3) ({
      %5 = "tf.OptionalGetValue"(%2) : (tensor<!tf_type.variant>) -> tensor<32xf32>
      %6 = func.call @__reduce_func_1(%arg3, %5) {_xla_compile_device_type = "TPU"} : (tensor<i64>, tensor<32xf32>) -> tensor<i64>
      "tf.Yield"(%6) : (tensor<i64>) -> ()

  let results = (outs
    TF_VariantTensor:$optional
  );

  TF_DerivedOperandTypeListAttr Toutput_types = TF_DerivedOperandTypeListAttr<0>;
}

def TF_OptionalGetValueOp : TF_Op<"OptionalGetValue", [Pure]> {
  let summary = [{
Returns the value stored in an Optional variant or raises an error if none exists.
  }];

  let arguments = (ins
    TF_VariantTensor:$optional
  );