public:
  LogicalResult matchAndRewrite(ConstantTensorOp cst_tensor_op,
                                PatternRewriter &rewriter) const override {
    Location loc = cst_tensor_op.getLoc();
    Type out_type = cst_tensor_op.getType();
    Operation *new_cst = nullptr;

    ArrayAttr array;
    if (matchPattern(cst_tensor_op.getArg(), m_Constant(&array))) {
      llvm::DenseSet<Type> all_types;

    ModuleOp module, llvm::ArrayRef<llvm::StringRef> function_names) {
  SymbolTableCollection symbol_table;
  OpBuilder builder(module.getContext());

  OwningOpRef<ModuleOp> pruned =
      builder.create<ModuleOp>(module->getLoc());
  (*pruned)->setAttrs(module->getAttrs());
  builder.setInsertionPointToEnd(pruned->getBody());

  llvm::SmallDenseSet<func::FuncOp> added;
  for (const llvm::StringRef function_name : function_names) {

  if (!llvm::hasSingleElement(func)) {
    LLVM_DEBUG(llvm::dbgs() << "Expect single block function, skip conversion "
                               "to tf_executor dialect\n");
    return;
  }
  auto loc = func.getLoc();
  mlir::Block& body = func.front();
  // Find region of interest and ReturnOp.
  auto copy_range = body.without_terminator();
  if (copy_range.begin() != copy_range.end() &&

          contraction,
          "cannot fuse with the subsequent op as it's not supported by the "
          "target device.");
    }

    SmallVector<Location, 3> locations{contraction.getLoc(), bias_add.getLoc()};
    SmallVector<Attribute, 2> fused_ops{StringAttr::get(
        context, bias_add.getOperation()->getName().stripDialect())};

    // BiasAdd may or may not feed into an activation function.

  // In some cases the shared_name attribute doesn't have the same
  // tensor name in the model, so we first try to use the location
  // then fallback to shared_name attribute.
  if (auto loc = mlir::dyn_cast<NameLoc>(var_handle_op->getLoc()))
    return loc.getName().str();
  return var_handle_op.getSharedName().str();
}

std::vector<std::string> GetVariableNames(
    llvm::ArrayRef<TF::VarHandleOp> var_handle_ops) {
  std::vector<std::string> names;

    ++dim;
  }

  auto permutation_and_shape = GetPermutationAndTransposedShape(
      permutation_array, updates_type, rewriter);

  auto transposed_updates = rewriter.create<mhlo::TransposeOp>(
      scatter_op->getLoc(), permutation_and_shape.shape, updates,
      permutation_and_shape.permutation);

  updates = transposed_updates;
  updates_type = permutation_and_shape.shape;
  return success();
}

}  // end namespace odml

};

template <typename OP>
int replace(OP op, int count) {
  OpBuilder builder(op);
  std::string name = absl::StrFormat("_iterator%d", count++);

  auto new_op = builder.create<TF::IteratorOp>(
      op->getLoc(), op->getResultTypes()[0], name, /*container=*/"",
      op.getOutputTypes(), op.getOutputShapes());
  op->getResults()[0].replaceAllUsesWith(new_op->getResults()[0]);
  if (op->use_empty()) op->erase();
  return count;
}

    IntegerAttr group_key = rewriter.getI32IntegerAttr(0);

    auto const_group_size = rewriter.create<TF::ConstOp>(
        assign_group->getLoc(), assign_group.getResult(0).getType(),
        group_size);
    auto const_group_key = rewriter.create<TF::ConstOp>(
        assign_group->getLoc(), assign_group.getResult(1).getType(), group_key);
    rewriter.replaceAllUsesWith(assign_group.getResult(0), const_group_size);

  }

  auto lhs_new_concat = rewriter.create<mhlo::ConcatenateOp>(
      concat->getLoc(), concat.getType().clone(lhs_new_concat_shape),
      all_dot_lhs, rewriter.getI64IntegerAttr(lhs_batch_dim));
  auto rhs_new_concat = rewriter.create<mhlo::ConcatenateOp>(
      concat->getLoc(), concat.getType().clone(rhs_new_concat_shape),
      all_dot_rhs, rewriter.getI64IntegerAttr(rhs_batch_dim));

  context.loadDialect<TF::TensorFlowDialect, quant::QuantizationDialect,
                      func::FuncDialect>();
  EmptyPatternRewriter pattern_rewriter(builder);
  Value value = CreateConstValue<float>(builder, module->getLoc(), {1024, 2},
                                        SmallVector<float>(2048, 0));

  QuantizationComponentSpec quant_spec;
  quant_spec.set_quantization_component(
      QuantizationComponentSpec::COMPONENT_WEIGHT);