Type result_type, func::FuncOp branch_func,
                            ConversionPatternRewriter *rewriter) const {
    auto guard = OpBuilder::InsertionGuard(*rewriter);
    auto inputs = branch_func.getFunctionType().getInputs();
    Block *block = rewriter->createBlock(
        &branch_func.getBody(), branch_func.begin(), inputs,

  void runOnOperation() override;
};

using mhlo::DotDimensionNumbersAttr;

// Replaces `region`'s terminator to TF::Yield.
void ReplaceReturnOp(Region& region, PatternRewriter& rewriter) {
  OpBuilder::InsertionGuard guard(rewriter);

  for (auto& block : region.getBlocks()) {
    Operation* terminator = block.getTerminator();
    auto return_op = llvm::dyn_cast_or_null<mhlo::ReturnOp>(terminator);