pm.addPass(mlir::quant::CreateUnfreezeConstantsPass());
      },
      ctx, module_op));

  if (const absl::Status create_dir_status =
          Env::Default()->CreateDir(std::string(checkpoint_dir));
      !create_dir_status.ok()) {
    LOG(ERROR) << "Failed to create checkpoint directory at: "
               << checkpoint_dir;
    return create_dir_status;
  }

  TF_ASSIGN_OR_RETURN(const auto unused_variable_names,

    return {0, 2, 3, 1};
  } else {
    return {};
  }
}

// Shuffle elements in the `attr` according to the permutation. Optional
// `inner_size` allows to shuffle array attributes created from rank 2 tensors
// on outer dimension only.
ArrayAttr ShuffleArrayAttr(ArrayAttr attr, ArrayRef<int64_t> permutation,
                           int inner_size) {
  if (attr.empty()) return attr;

    // If the op is rewritten, FusedLoc can be created.
    for (Location child_loc : mlir::cast<FusedLoc>(loc).getLocations()) {
      std::optional<QuantizationUnitLoc::QuantizationUnit> found_unit =
          FindQuantizationUnitFromLoc(child_loc);
      if (found_unit.has_value()) return found_unit;
    }
  } else if (isa<CallSiteLoc>(loc)) {
    // If the graph is inlined, CallSiteLoc can be created.
    return FindQuantizationUnitFromLoc(

constexpr char kXlaOutsideCompilationAttr[] = "_xla_outside_compilation";

bool HasOutsideCompilationAttribute(Operation* op) {
  return op->getAttrOfType<StringAttr>(kXlaOutsideCompilationAttr) != nullptr;
}

// Finds op that created a given value. If the value is a BlockArgument, this
// returns the owner of the Block.
Operation* GetOpOfValue(Value value) {
  if (auto block_arg = mlir::dyn_cast<BlockArgument>(value))

  const std::string dump_file_path =
      tsl::io::JoinPath(*dump_dir, dump_file_name);
  TF_ASSIGN_OR_RETURN(std::unique_ptr<llvm::raw_ostream> file,
                      WritableFileWrapper::Create(dump_file_path));

  LOG(INFO) << "IR dump file created: " << dump_file_path;
  return file;
}

class PrinterConfig : public mlir::PassManager::IRPrinterConfig {
 public:
  explicit PrinterConfig(

      report_file_path));

  const LogicalResult run_result = pm.run(*module_op);
  ASSERT_TRUE(succeeded(run_result));

  // The report file is not created because `QuantizeCompositeFunctionsPass` was
  // not run.
  EXPECT_THAT(ReadFileToString(report_file_path),
              StatusIs(absl::StatusCode::kNotFound));
}

TEST_F(SaveQuantizationReportInstrumentationTest,

  auto new_launch_op = builder.create<tf_device::LaunchOp>(
      old_launch_op->getLoc(), old_launch_op->getDeviceAttr(),
      /*result_types=*/new_launch_op_results_types);

  new_launch_op.getBody().takeBody(old_launch_op->getBody());
  new_launch_op.GetBody().getTerminator()->setOperands(new_launch_op_results);

  return new_launch_op;
}

// Create the new device launch op which wraps the copy op.

                                                new_results)))
      return failure();

    // Update the regions. The dialect conversion framework wants new regions to
    // be created and updated, rather than updating the old op. Thus we use an
    // OperationState so we can add regions to the new op.
    OperationState state(op->getLoc(), op->getName().getStringRef(), operands,

      OpBuilder replacer(yield_op);
      island_op.getBody().push_back(new Block);
      replacer.create<mlir::func::ReturnOp>(yield_op.getLoc(),
                                            yield_op.getOperands());
      yield_op.erase();
    }

    // Remap the captured operands in the (former) island block with newly
    // created entry block arguments in the function body.
    {

  ConvertAssetArgsTest() {
    ctx_.loadDialect<func::FuncDialect, TF::TensorFlowDialect,
                     tf_saved_model::TensorFlowSavedModelDialect>();
  }

  // Parses `module_op_str` to create a `ModuleOp`. Checks whether the created
  // module op is valid.
  OwningOpRef<ModuleOp> ParseModuleOpString(
      const absl::string_view module_op_str) {
    auto module_op_ref = parseSourceString<ModuleOp>(module_op_str, &ctx_);