// have device attribute.
        op->setAttr(kDeviceAttr, builder.getStringAttr(default_device_));
      }
    });
  }
};

#define GEN_PASS_DEF_TFDEVICEASSIGNMENTBYFUNCATTRPASS
#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_passes.h.inc"

// A pass to perform device assignment for TF dialect ops that do not

  // Note: This module search for cluster ops is a bit wasteful as we could have
  // collected many cluster ops when we were populating reachable functions. But
  // we would still need to do a walk to find all clusters that do not
  // reference any function.
  for (func::FuncOp func : module.getOps<func::FuncOp>()) {
    // If we have already applied patterns to a function then we can skip

  const auto num_replicas = replicate_op.getN();
  auto return_op = llvm::dyn_cast<tf_device::ReturnOp>(
      replicate_op.getRegion().front().getTerminator());

  // Get the new result types.
  // TODO(prakalps): Do not add a value to returned values if it is already
  // returned.
  auto new_result_types = llvm::to_vector<4>(replicate_op->getResultTypes());
  for (auto assign : tail_assign_variable_ops) {

  mlir::Builder builder(ctx);

  // Collect devices with attached metadata.
  llvm::SmallVector<mlir::NamedAttribute, 8> devices;
  devices.reserve(device_set->devices().size());

  // For device that do not have any metadata, or if we failed to parse metadata
  // from the DeviceSet, we add a unit attribute to the `tf.devices` attribute.
  for (Device* device : device_set->devices()) {

    }

    auto layout = mlir::GetTPUInfeedLayout(
        mlir::TupleType::get(builder.getContext(), result_types), builder);
    if (failed(layout)) return mlir::WalkResult::interrupt();
    // Do not append a UnitAttr for the "token" operand here to avoid
    // compilation failure when exporting the "layouts" attribute to a graph
    // node. Instead, add the UnitAttr during LegalizeTF pass.

  if (options_.legalize_to_tflite_ops) {
    // After we creat the alternative subgraph, we can still do canonicalization
    // legalization & other optimizations as long as we're not inlining the
    // function.
    // And in fact, we probably need to do the proper legalization, for the
    // compute cost to work. (in case we added some TF ops)
    pass_manager->addPass(mlir::TFL::CreatePrepareTFPass(

// pathological cases haven't hit us yet.
//
// 2. Cloning functions can help by allowing code to be specialized (much as
// inlining does). In fact, tf-shape-inference attempts to do specialization
// of callees which is difficult if callees have multiple uses.
class GuaranteeAllFuncsOneUse
    : public impl::GuaranteeAllFuncsOneUsePassBase<GuaranteeAllFuncsOneUse> {
 public:

  if (absl::Status status = AddCalibrationStatistics(
          module_op, calibration_data_dir, config.calibration_options(),
          *py_function_lib_);
      !status.ok()) {
    LOG(WARNING) << "Some CustomAggregator ops do not have min or max "
                    "values. Parts of the graph are not quantized. "
                 << status;
  }

  return module_op;
}

    builder.create<TF::AssignVariableOp>(
        restore_op.getLoc(), target_var_handle_ops[idx], restore_result);
  }
}

// TODO(b/261813194): Do not create a new RestoreV2 op when a RestoreV2 op
// already exists.
void InsertRestoreOpPass::runOnOperation() {
  ModuleOp module_op = getOperation();

  func::FuncOp session_init_func = GetInitializerFunction(

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

// This file wraps TensorFlow Graph(Def) to MLIR module conversion into passes
// to satisfy the API of MLIR pass registration. In order to do this, the
// command-line option header is pulled in.

#include <memory>
#include <utility>

#include "absl/container/flat_hash_set.h"
#include "mlir/IR/BuiltinOps.h"  // from @llvm-project