if (broadcasted_shape != result_type.getShape()) continue;

    // Update the operand of the op to be the operand of the broadcast.
    rewriter.modifyOpInPlace(
        op, [&]() { op->getOpOperand(i).set(broadcast.getInput()); });
    changed = true;
  }
  return success(changed);
}

void BroadcastFoldPass::runOnOperation() {

    }
    auto newType =
        FunctionType::get(rewriter.getContext(), argument_types, return_types);
    if (f.getFunctionType() == newType) {
      return failure();
    }
    rewriter.modifyOpInPlace(f, [&] { f.setType(newType); });

    // Adjust the type of the return values callers of the function to
    // match the "func.return" within the function.
    //
    // So this would transform

  IRRewriter rewriter(op->getContext());

  op->walk([&](Operation *op) {
    if (op->getDialect()->getNamespace() != "vhlo") return;

    // Convert operands
    rewriter.modifyOpInPlace(op, [&]() {
      rewriter.setInsertionPoint(op);
      WrapOperandsInUnrealizedCastAndConvert(op, converter, rewriter);

      // Convert op types
      for (auto value : op->getResults()) {

  using OpConversionPattern::OpConversionPattern;

  LogicalResult matchAndRewrite(
      func::ReturnOp op, OpAdaptor adaptor,
      ConversionPatternRewriter &rewriter) const override {
    rewriter.modifyOpInPlace(op,
                             [&] { op->setOperands(adaptor.getOperands()); });
    return success();
  }
};

struct ConvertYield : public OpConversionPattern<TF::YieldOp> {