"TFL::IsReducedTailOfShape($0.getType(), $1.getType())">>;

def IsRankLessThanEqualTo : Constraint<CPred<
  "$0.getType().cast<ShapedType>().getRank() <= "
  "$1.getType().cast<ShapedType>().getRank()">>;

def Flatten : NativeCodeCall<
  "$0.cast<DenseElementsAttr>()"
    ".reshape(RankedTensorType::get({$0.getType().cast<ShapedType>().getNumElements()}, "
                                   "$0.getType().cast<ShapedType>().getElementType()))">;

      rhs,
      DenseIntElementsAttr::get(
          RankedTensorType::get({rhs_rank}, rewriter.getI64Type()),
          rhs_permutation));
  // Reshapes lhs to flatten out_dimensions and contracting_dimensions.
  llvm::SmallVector<int64_t, 4> lhs_flattened_shape = Concat<int64_t>(
      lhs_dot_dimensions_info.batch_dimensions().SizesArray(),
      llvm::ArrayRef<int64_t>{

// TODO(b/190203981): Move to a separate nest-like library.
void Flatten(const TaggedValue& value,
             std::vector<AbstractTensorHandle*>* flat_args) {
  if (value.type() == TaggedValue::Type::TENSOR) {
    flat_args->emplace_back(value.tensor().get());
  } else if (value.type() == TaggedValue::Type::TUPLE) {
    for (const auto& t : value.tuple()) {
      Flatten(t, flat_args);
    }
  } else {

    llvm::StringRef equation, RankedTensorType lhs_ty) {
  llvm::StringRef lhs;
  llvm::StringRef out;
  std::tie(lhs, out) = equation.split("->");
  if (lhs.empty() || out.empty()) return std::nullopt;

  // Try to flatten the "..." if possible.
  int lhs_named_label, rhs_named_label;

  // following rhs and rhs_ty variables are non-functional here only created to
  // comply with the existing API
  llvm::StringRef rhs;

         (std::equal(i1, reduced_e1, i2));
}

// Check if the value of the last dimension of type1 is equal to the number of
// elements in type2. This is a required condition to flatten type2 to form a
// 1D array and allow the binaryOp handle the broadcasting implicitly.
bool IsLastDimEqualToNumElements(Type type1, Type type2) {
  return (mlir::cast<ShapedType>(type1).getRank() >= 1 &&

// linearized function argument or return on a FunctionDef, and hence to an
// mlir::func::FuncOp argument / return.
//
// This must match the linearization that happens in `tf.nest.flatten`.
// In particular, dict values should be linearized in sorted key order.
//
// The linearized index paths can be returned back to a structured
// representation (e.g. to emit C structs matching a signature) with a simple

@chanis, Chenyang Liu, Corey Wharton, Daeyun Shin, Daniel Julius Lasiman, Daniel
Waterworth, Danijar Hafner, Darren Garvey, Denis Gorbachev, @DjangoPeng,
Egor-Krivov, Elia Palme, Eric Platon, Fabrizio Milo, Gaetan Semet, Georg
Nebehay, Gu Wang, Gustav Larsson, @haosdent, Harold Cooper, Hw-Zz, @ichuang,
Igor Babuschkin, Igor Macedo Quintanilha, Ilya Edrenkin, @ironhead, Jakub

  // Add inline pass.
  pass_manager.addPass(mlir::createInlinerPass());

  // Expands mhlo.tuple ops.
  pass_manager.addPass(
      mlir::mhlo::createExpandHloTuplesPass(entry_function_name.str()));
  // Flatten tuples for control flows.
  pass_manager.addNestedPass<mlir::func::FuncOp>(
      mlir::mhlo::createFlattenTuplePass());

  mlir::odml::AddMhloOptimizationPasses(
      pass_manager,

      rhs,
      DenseIntElementsAttr::get(
          RankedTensorType::get({rhs_rank}, rewriter.getI64Type()),
          rhs_permutation));

  // Reshapes lhs to flatten out_dimensions and contracting_dimensions.
  llvm::SmallVector<int64_t, 4> lhs_flattened_shape = Concat<int64_t>(
      lhs_dot_dimensions_info.batch_dimensions().SizesArray(),
      llvm::ArrayRef<int64_t>{

}

// Determine if op commutes with transposes. Requires a strict
// definition of Elementwise, all i/o shapes and types must be same-rank
// broadcastable and fully static. Consider moving this into attribute later.
bool IsElementwise(Operation *op) {
  if (!(llvm::isa<TFL::AddOp, TFL::MulOp, TFL::DivOp, TFL::SubOp,
                  TFL::MaximumOp, TFL::MinimumOp>(op))) {
    return false;
  }

  auto opr1_type =