rewriter.replaceOpWithNewOp<mhlo::ReturnOp>(op, adaptor.getOperands());
    return success();
  }
};

// Returns a new tensor type from the given type with element type updated to
// the given type.
TensorType UpdateElementTypeTo(Type ty, Type element_ty) {
  auto ranked_ty = mlir::dyn_cast<RankedTensorType>(ty);
  if (!ranked_ty) {
    return UnrankedTensorType::get(element_ty);
  }