"Const",
           {},
           {{"value", test::AsScalar<int32>(2)}, {"dtype", DT_INT64}}},
          {{"scale"},
           "Cast",
           {"two"},
           {{"SrcT", DT_INT64}, {"DstT", DT_FLOAT}}},
          {{"y"}, "Mul", {"x", "scale"}, {{"T", DT_FLOAT}}},
      });

  tensorflow::set_tf2_execution(true);
  FunctionDefLibrary flib;
  *flib.add_function() = fd;

  if (node_def->op() == "Cast" && attr->contains("SrcT") &&
      attr->contains("DstT") &&
      attr->at("SrcT").type() == attr->at("DstT").type() &&
      attr->contains("Truncate") && !attr->at("Truncate").b()) {
    node_def->set_op("Identity");
    attr->insert({{"T", attr->at("SrcT")}});
    attr->erase("SrcT");
    attr->erase("DstT");
    attr->erase("Truncate");
  }
}

}  // namespace

      // Attr def
      {"T: {float, double, int32, int64}"},
      // Nodes
      {
          {{"two"}, "Const", {}, {{"value", kTwo}, {"dtype", DT_INT64}}},
          {{"scale"}, "Cast", {"two"}, {{"SrcT", DT_INT64}, {"DstT", "$T"}}},
          {{"y"}, "Mul", {"x", "scale"}, {{"T", "$T"}}},
      });
}

class MlirGraphOptimizationPassTest : public Test {
 public:
  void Init(Status pass_run_result,

  let summary = "Cast x of type SrcT to y of DstT.";

  let arguments = (ins
    TF_Tensor:$x,

    DefaultValuedOptionalAttr<BoolAttr, "false">:$Truncate
  );

  let results = (outs
    TF_Tensor:$y
  );

  TF_DerivedOperandTypeAttr SrcT = TF_DerivedOperandTypeAttr<0>;
  TF_DerivedResultTypeAttr DstT = TF_DerivedResultTypeAttr<0>;

  let hasFolder = 1;
}