Type result_elt_type = GetFixedElementType(tfr_type_attr, rewriter);
    if (!result_elt_type) {
      return cast_op.getArg();
    }

    Type original_input_type =
        mlir::cast<TypeAttr>(cast_op.getInputElementType()).getValue();
    if (result_elt_type != original_input_type) {
      UnrankedTensorType result_type = UnrankedTensorType::get(result_elt_type);
      return rewriter.create<TF::CastOp>(loc, result_type, cast_op.getArg());

  let extraClassDeclaration = [{
    // Return element type of the input tensor type. Only available when the
    // input is a MLIR built-in tensor type.
    Attribute getInputElementType() {
      if (auto ty = getArg().getType().dyn_cast<TensorType>()) {
        return TypeAttr::get(ty.getElementType());
      }
      return {};
    }
  }];

  let hasCanonicalizer = 1;
}