for (const BlockArgument &arg : entry.getArguments()) {
    signature_conversion.addInputs(arg.getArgNumber(),
                                   updated_argument_types[arg.getArgNumber()]);
  }
  rewriter.applySignatureConversion(&entry.front(), signature_conversion);
}

// Changes the function type of `cond_func` and `body_func` for the given While
// op.
LogicalResult UpdateFunctionTypesForWhileOp(

          TensorType updated_ty =
              UpdateElementTypeTo(original_ty, element_types[block_idx]);
          signature.addInputs(block_idx, {updated_ty});
        }
        rewriter.applySignatureConversion(&region.front(), signature);
      }
    }

    // Replace all uses of `op` results with the newly created op.
    rewriter.replaceOp(op, mhlo_op);
    return success();
  }
};