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());

    return success();
  }
};

quant::QuantizedType getQuantizedElementType(CastOp cast_op) {
  if (!cast_op || !cast_op.getInputElementType()) {
    return {};
  }
  return cast_op.getInputElementType()
      .cast<TypeAttr>()
      .getValue()
      .dyn_cast<quant::QuantizedType>();
}

class RemoveRawDataOp : public OpRewritePattern<TFRQuantRawDataOp> {

  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;
}