code: #include <torch/csrc/jit/tensorexpr/ir_printer.h>
  
  #include <torch/csrc/jit/tensorexpr/ir_simplifier.h>
  
  namespace torch {
  namespace jit {
  namespace tensorexpr {
  
  void IRPrinter::print(ExprHandle expr) {
    expr.node()->accept(this);
  }
  
  void IRPrinter::print(const Expr& expr) {
    expr.accept(this);
  }
  
  void IRPrinter::print(const Stmt& stmt) {
    stmt.accept(this);
  }
  
  // TODO: change whether to include the parenthesis to the parent expression,
  // we need to look at the operator precedence to make the output simpler.
  template <typename Op>
  void visitBinaryOp(
      const BinaryOpNode<Op>* v,
      const std::string& op_str,
      IRPrinter* printer,
      bool parens = true) {
    std::ostream& os = printer->os();
    int self_prec = getPrecedence(v->expr_type());
    int lhs_prec = getPrecedence(v->lhs()->expr_type());
    int rhs_prec = getPrecedence(v->rhs()->expr_type());
  
    if (lhs_prec >= self_prec) {
      os << "(";
    }
    v->lhs()->accept(printer);
    if (lhs_prec >= self_prec) {
      os << ")";
    }
  
    os << " " << op_str << " ";
  
    if (rhs_prec >= self_prec) {
      os << "(";
    }
    v->rhs()->accept(printer);
    if (rhs_prec >= self_prec) {
      os << ")";
    }
  }
  
  void IRPrinter::visit(const Add* v) {
    visitBinaryOp(v, "+", this);
  }
  
  void IRPrinter::visit(const Sub* v) {
    visitBinaryOp(v, "-", this);
  }
  
  void IRPrinter::visit(const Mul* v) {
    visitBinaryOp(v, "*", this);
  }
  
  void IRPrinter::visit(const Div* v) {
    visitBinaryOp(v, "/", this);
  }
  
  void IRPrinter::visit(const And* v) {
    visitBinaryOp(v, "&", this);
  }
  
  void IRPrinter::visit(const Or* v) {
    visitBinaryOp(v, "|", this);
  }
  
  void IRPrinter::visit(const Xor* v) {
    visitBinaryOp(v, "^", this);
  }
  
  void IRPrinter::visit(const Lshift* v) {
    visitBinaryOp(v, "<<", this);
  }
  
  void IRPrinter::visit(const Rshift* v) {
    visitBinaryOp(v, ">>", this);
  }
  
  void IRPrinter::visit(const Mod* v) {
    if (v->dtype().is_integral()) {
      visitBinaryOp(v, "%", this);
    } else if (v->dtype().is_floating_point()) {
      os() << "mod(" << v->lhs() << ", " << v->rhs() << ")";
    } else {
      throw std::runtime_error("invalid dtype: " + std::to_string(v->dtype()));
    }
  }
  
  void IRPrinter::visit(const Max* v) {
    os() << "Max(";
    v->lhs()->accept(this);
    os() << ", ";
    v->rhs()->accept(this);
    os() << ", " << (unsigned int)v->propagate_nans() << ")";
  }
  
  void IRPrinter::visit(const Min* v) {
    os() << "Min(";
    v->lhs()->accept(this);
    os() << ", ";
    v->rhs()->accept(this);
    os() << ", " << (unsigned int)v->propagate_nans() << ")";
  }
  
  void IRPrinter::visit(const CompareSelect* v) {
    CompareSelectOperation cmp_op = v->compare_select_op();
    int self_prec = getPrecedence(v->expr_type());
    int lhs_prec = getPrecedence(v->lhs()->expr_type());
    int rhs_prec = getPrecedence(v->rhs()->expr_type());
  
    if (lhs_prec >= self_prec) {
      os() << "(";
    }
    v->lhs()->accept(this);
    if (lhs_prec >= self_prec) {
      os() << ")";
    }
    switch (cmp_op) {
      case CompareSelectOperation::kEQ:
        os() << "==";
        break;
      case CompareSelectOperation::kNE:
        os() << "!=";
        break;
      case CompareSelectOperation::kGT:
        os() << ">";
        break;
      case CompareSelectOperation::kGE:
        os() << ">=";
        break;
      case CompareSelectOperation::kLT:
        os() << "<";
        break;
      case CompareSelectOperation::kLE:
        os() << "<=";
        break;
      default:
        throw std::runtime_error("invalid compare select operator");
    }
  
    if (rhs_prec >= self_prec) {
      os() << "(";
    }
    v->rhs()->accept(this);
    if (rhs_prec >= self_prec) {
      os() << ")";
    }
    os() << " ? ";
  
    auto withParens = [&](const Expr* e) {
      auto prec = getPrecedence(e->expr_type());
      if (prec >= self_prec) {
        os() << "(";
      }
      e->accept(this);
      if (prec >= self_prec) {
        os() << "(";
      }
    };
    withParens(v->ret_val1());
    os() << " : ";
    withParens(v->ret_val2());
  }
  
  static void formatFPSuffix(std::ostream& os, double v) {
    // No suffix for doubles.
  }
  
  template <typename T>
  static void formatFPSuffix(std::ostream& os, T v) {
    os << (v == std::ceil(v) ? ".f" : "f");
  }
  
  template <
      typename T,
      std::enable_if_t<std::is_floating_point<T>::value>* = nullptr>
  static void formatImm(std::ostream& os, T v) {
    const int precision = 16;
    if (std::isnan(v)) {
      os << "NAN";
    } else if (std::isinf(v)) {
      os << (v > 0 ? "POS_INFINITY" : "NEG_INFINITY");
    } else {
      os << std::setprecision(precision) << v;
      formatFPSuffix(os, v);
    }
  }
  
  template <
      typename T,
      std::enable_if_t<!std::is_floating_point<T>::value>* = nullptr>
  static void formatImm(std::ostream& os, T v) {
    os << +v;
  }
  
  // NOLINTNEXTLINE
  #define IMM_PRINT_VISIT(Type, Name)           \
    void IRPrinter::visit(const Name##Imm* v) { \
      formatImm(os(), v->value());              \
    }
  AT_FORALL_SCALAR_TYPES_AND2(Bool, Half, IMM_PRINT_VISIT);
  #undef IMM_PRINT_VISIT
  
  void IRPrinter::visit(const Cast* v) {
    auto dtype = v->dtype();
    os() << dtype << "(";
    v->src_value()->accept(this);
    os() << ")";
  }
  
  void IRPrinter::visit(const Var* v) {
    os() << name_manager_.get_unique_name(v);
  }
  
  void IRPrinter::visit(const Let* v) {
    int self_prec = getPrecedence(v->expr_type());
    int value_prec = getPrecedence(v->value()->expr_type());
    int body_prec = getPrecedence(v->body()->expr_type());
    os() << "let ";
    v->var()->accept(this);
    os() << " = ";
  
    if (value_prec >= self_prec) {
      os() << "(";
    }
    v->value()->accept(this);
    if (value_prec >= self_prec) {
      os() << ")";
    }
  
    os() << " in ";
  
    if (body_prec >= self_prec) {
      os() << "(";
    }
    v->body()->accept(this);
    if (body_prec >= self_prec) {
      os() << ")";
    }
  }
  
  void IRPrinter::visit(const LetStmt* v) {
    const Var* var = v->var();
    os() << var->dtype().ToCppString() << " " << *var << " = " << *v->value()
         << "; " << std::endl;
    v->body()->accept(this);
  }
  
  void IRPrinter::visit(const Ramp* v) {
    os() << "Ramp(" << *v->base() << ", " << *v->stride() << ", " << v->lanes()
         << ")";
  }
  
  void IRPrinter::visit(const Load* v) {
    // TODO: support the mask case
    os() << *v->base_handle() << "[" << *v->index() << "]";
  }
  
  void IRPrinter::visit(const For* v) {
    const Var* var = v->var();
    VarHandle vv(var);
    emitIndent();
    os() << "for (" << var->dtype().ToCppString() << " " << vv << " = "
         << ExprHandle(v->start()) << "; " << vv << " < " << ExprHandle(v->stop())
         << "; " << vv << "++) {";
    std::string loop_options_str = v->loop_options().ToString();
    if (!loop_options_str.empty()) {
      os() << " // " << loop_options_str;
    }
    os() << std::endl;
    if (v->body()) {
      indent_++;
      os() << *v->body();
      indent_--;
    }
    emitIndent();
    os() << "}";
  }
  
  void IRPrinter::visit(const Block* v) {
    for (Stmt* s : v->stmts()) {
      os() << *s << std::endl;
    }
  }
  
  void IRPrinter::visit(const Store* v) {
    // TODO: handle the mask
    emitIndent();
    os() << *v->base_handle() << "[" << *v->index() << "] = " << *v->value()
         << ";";
  }
  
  void IRPrinter::visit(const Broadcast* v) {
    os() << "Broadcast(" << *v->value() << ", " << v->lanes() << ")";
  }
  
  void IRPrinter::visit(const IfThenElse* v) {
    os() << "IfThenElse(" << *v->condition() << ", " << *v->true_value() << ", "
         << *v->false_value() << ")";
  }
  
  void IRPrinter::visit(const BaseCallNode* v) {
    os() << v->func_name() << "(";
    for (int i = 0; i < v->nparams(); i++) {
      if (i > 0) {
        os() << ", ";
      }
      os() << *v->param(i);
    }
    os() << ")";
  }
  
  void IRPrinter::visit(const Allocate* v) {
    emitIndent();
    os() << "Allocate(" << *v->buffer_var() << ", " << v->dtype();
    os() << ", {";
    const std::vector<const Expr*>& dims = v->dims();
    for (size_t i = 0; i < dims.size(); i++) {
      if (i != 0) {
        os() << ", ";
      }
      os() << *dims[i];
    }
    os() << "});";
  }
  
  void IRPrinter::visit(const Free* v) {
    emitIndent();
    os() << "Free(" << *v->buffer_var() << ");";
  }
  
  void IRPrinter::visit(const Cond* v) {
    const Expr* cond = v->condition();
    Stmt* true_stmt = v->true_stmt();
    Stmt* false_stmt = v->false_stmt();
    if (!true_stmt) {
      emitIndent();
      os() << "if (!" << *cond << ") {" << std::endl;
      indent_++;
      os() << *false_stmt << std::endl;
      indent_--;
      emitIndent();
      os() << "}";
    } else {
      emitIndent();
      os() << "if (" << *cond << ") {" << std::endl;
      indent_++;
      os() << *true_stmt << std::endl;
      indent_--;
      emitIndent();
      os() << "}";
      if (false_stmt) {
        os() << " else {" << std::endl;
        indent_++;
        os() << *false_stmt << std::endl;
        indent_--;
        emitIndent();
        os() << "}";
      }
    }
  }
  
  void IRPrinter::visit(const Term* v) {
    os() << "Term(";
    v->scalar()->accept(this);
    for (auto* t : v->variables()) {
      os() << ",";
      t->accept(this);
    }
    os() << ")";
  }
  
  void IRPrinter::visit(const Polynomial* v) {
    bool first = true;
    os() << "Polynomial(";
    for (auto* t : v->variables()) {
      emitIndent();
      if (!first) {
        os() << " + ";
      }
      first = false;
      t->accept(this);
    }
  
    if (!first) {
      os() << " + ";
    }
    v->scalar()->accept(this);
    os() << ")";
  }
  
  void IRPrinter::emitIndent() {
    os() << std::setw(2 * indent_) << "";
  }
  
  std::ostream& operator<<(std::ostream& stream, const ExprHandle& expr) {
    IRPrinter::PrinterStream* printer_stream =
        dynamic_cast<IRPrinter::PrinterStream*>(&stream);
    if (printer_stream != nullptr) {
      expr.node()->accept(printer_stream->printer());
    } else {
      IRPrinter p(stream);
      p.print(expr);
    }
    return stream;
  }
  
  std::ostream& operator<<(std::ostream& stream, const Expr& expr) {
    IRPrinter::PrinterStream* printer_stream =
        dynamic_cast<IRPrinter::PrinterStream*>(&stream);
    if (printer_stream != nullptr) {
      expr.accept(printer_stream->printer());
    } else {
      IRPrinter p(stream);
      p.print(expr);
    }
    return stream;
  }
  
  std::ostream& operator<<(std::ostream& stream, const Stmt& stmt) {
    IRPrinter::PrinterStream* printer_stream =
        dynamic_cast<IRPrinter::PrinterStream*>(&stream);
    if (printer_stream != nullptr) {
      stmt.accept(printer_stream->printer());
    } else {
      IRPrinter p(stream);
      p.print(stmt);
    }
    return stream;
  }
  
  void print(const Expr* expr) {
    if (expr) {
      IRPrinter p(std::cout);
      p.print(*expr);
    } else {
      std::cout << "(null expr)";
    }
  }
  
  void print(const Stmt* stmt) {
    if (stmt) {
      IRPrinter p(std::cout);
      p.print(*stmt);
    } else {
      std::cout << "(null stmt)\n";
    }
  }
  
  } // namespace tensorexpr
  } // namespace jit
  } // namespace torch
  
  namespace std {
  std::string to_string(const Expr* expr) {
    std::ostringstream oss;
    oss << *expr;
    return oss.str();
  }
  
  std::string to_string(const Stmt* stmt) {
    std::ostringstream oss;
    oss << *stmt;
    return oss.str();
  }
  } // namespace std
  