auto attr = node_def->mutable_attr();
  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");
  }
}

          {{"two"},
           "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;

  device: "/job:localhost/replica:0/task:0/device:GPU:0"
  attr {
    key: "DstT"
    value {
      type: DT_HALF
    }
  }
  attr {
    key: "SrcT"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "Truncate"
    value {
      b: false
    }
  }
}
node {

      {"y: T"},
      // 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,

  device: "/job:localhost/replica:0/task:0/device:GPU:0"
  attr {
    key: "DstT"
    value {
      type: DT_INT64
    }
  }
  attr {
    key: "SrcT"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "Truncate"
    value {
      b: false
    }
  }
}
node {

def TF_CastOp : TF_Op<"Cast", [Pure, SameOperandsAndResultShape]> {
  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>;