#include "tensorflow/compiler/mlir/tfr/integration/tfr_decompose_ctx.h"

#include <string>
#include <vector>

#include "absl/strings/str_cat.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/MemoryBuffer.h"
#include "llvm/Support/SMLoc.h"
#include "llvm/Support/SourceMgr.h"
#include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project

    if (!consumed_ops.contains(&op) && !IsConstant(&op)) {
      op.erase();
    }
  body.eraseArguments(erase_indices);
  cond.eraseArguments(erase_indices);

  llvm::SmallVector<Value> new_operands;
  llvm::SmallVector<Type> new_result_types;
  new_operands.reserve(n - erase_indices.size());
  new_result_types.reserve(n - erase_indices.size());

==============================================================================*/

#include "tensorflow/compiler/mlir/tf2xla/internal/clustering_bridge_passes.h"

#include <string>

#include "absl/log/log.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
#include "mlir/Pass/PassManager.h"  // from @llvm-project
#include "mlir/Transforms/Passes.h"  // from @llvm-project

#include <memory>
#include <string>
#include <utility>

#include "absl/status/status.h"
#include "absl/strings/string_view.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
#include "mlir/IR/Attributes.h"  // from @llvm-project
#include "mlir/IR/Builders.h"  // from @llvm-project

}

FailureOr<Attribute> GetTPUInfeedLayout(const ArrayRef<Type> types,
                                        OpBuilder &rewriter) {
  auto i64_type = rewriter.getIntegerType(64);
  if (types.size() > 1) {
    llvm::SmallVector<mlir::Attribute> v;
    v.reserve(types.size());
    for (const mlir::Type &t : types) {
      if (mlir::isa<mhlo::TokenType>(t)) continue;
      auto layout = GetTPUInfeedLayout({t}, rewriter);

  if (parallel_execute.getRegions().size() == 1) {
    builder->setInsertionPoint(parallel_execute);
    auto& block = parallel_execute.GetRegionBlockWithIndex(0);
    llvm::SmallVector<Operation*, 2> ops_move;
    for (Operation& op : block) {
      ops_move.push_back(&op);
    }
    if (ops_move.size() != 2) {
      parallel_execute.emitError() << "Expected 2 ops in parallel_execute.";

// tensor.
DenseFPElementsAttr GetCorrectionMatrix(Builder& builder, CompositeOp op) {
  const TorchAvgPoolData pool = GetTorchAvgPoolData(op);

  llvm::SmallVector<int64_t, 4> nhwc_out_shape(4);
  nhwc_out_shape[0] = 1;  // Broadcast batch.
  nhwc_out_shape[1] = pool.h_out;
  nhwc_out_shape[2] = pool.w_out;
  nhwc_out_shape[3] = 1;  // Broadcast channels.

// in `tensorflow/compiler/mlir/tensorflow/translate/import_model.cc`.
void FindQuantizationUnitsRecursively(Location loc,
                                      SmallVector<QuantizationUnit>& units) {
  if (!isa<FusedLoc>(loc)) return;

  auto set_node_and_func_name = [](QuantizationUnit& new_unit,
                                   StringRef name_loc_id) {
    if (name_loc_id.contains("@")) {

See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "llvm/ADT/SmallVector.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/Debug.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
#include "mlir/IR/Operation.h"  // from @llvm-project

#include "absl/base/casts.h"
#include "absl/container/inlined_vector.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "llvm/ADT/APFloat.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/SmallVector.h"
#include "mlir/IR/Attributes.h"  // from @llvm-project
#include "mlir/IR/Builders.h"  // from @llvm-project
#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project