namespace {
struct MarkInputOutputAliasesPass
    : public impl::MarkInputOutputAliasesPassBase<MarkInputOutputAliasesPass> {
  void runOnOperation() override;
};

constexpr char kAliasingAttr[] = "tf.aliasing_output";
constexpr int kUnassigned = -1;

struct AliasInfo {
  AliasInfo() : input_index(kUnassigned), output_index(kUnassigned) {}
  int input_index;
  int output_index;
};

InitMlir::InitMlir(int *argc, char ***argv) {
  llvm::setBugReportMsg(
      "TensorFlow crashed, please file a bug on "
      "https://github.com/tensorflow/tensorflow/issues with the trace "
      "below.\n");

  constexpr char kSeparator[] = "--";

  // Find index of separator between two sets of flags.
  int pass_remainder = 1;
  bool split = false;
  for (int i = 0; i < *argc; ++i) {
    if (llvm::StringRef((*argv)[i]) == kSeparator) {

#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"
#include "tensorflow/core/util/device_name_utils.h"

namespace mlir {
namespace TFTPU {

namespace {

constexpr char kDeviceAttr[] = "device";
constexpr char kDeviceCPU[] = "CPU";
constexpr char kFuncDeviceAttr[] = "tf.device";

struct TPUDynamicLayoutPass
    : public TF::PerFunctionAggregateAnalysisConsumerPass<

  CreateTFStandardPipeline(bridge, pipeline_options);

  mlir::StatusScopedDiagnosticHandler diag_handler(
      module.getContext(), /*propagate=*/false,
      /*filter_stack=*/!VLOG_IS_ON(1));

  constexpr char kBridgeComponent[] = "TFXLABridge";
  if (enable_logging || VLOG_IS_ON(1)) {
    tensorflow::DumpMlirOpToFile(kStandardPipelineBefore, module, "", &bridge);
    if (VLOG_IS_ON(2)) {

using DeviceMap = DenseMap<StringAttr, int64_t>;
constexpr int64_t kAnyDevice = 0;
constexpr ResourceAndDevice kInvalidResourceAndDevice{kInvalidResourceId,
                                                      kAnyDevice};
constexpr ResourceAndDevice kUnknownResourceAndDevice{kUnknownResourceId,
                                                      kAnyDevice};

constexpr char kDeviceAttr[] = "device";

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"

namespace mlir {
namespace TF {

namespace {

constexpr char kShapeInvariantAttr[] = "shape_invariant";

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

class DropWhileShapeInvariantPass

using mlir::RankedTensorType;
using mlir::StringAttr;
using mlir::Value;
using mlir::WalkResult;

constexpr char kReplicateSharding[] = "";
constexpr char kShardingAttr[] = "mhlo.sharding";
constexpr char kUseSpmdAttr[] = "use_spmd_for_xla_partitioning";
constexpr char kAliasingAttr[] = "tf.aliasing_output";
constexpr char kNumCoresPerReplicaAttr[] = "num_cores_per_replica";

#define GEN_PASS_DEF_TPUSHARDINGIDENTIFICATIONPASS

// Device coordinates are defined as (x, y, z, core), thus resulting in a rank 4
// topology.
constexpr int kTPUTopologyRank = 4;

constexpr char kDeviceTPUSystem[] = "TPU_SYSTEM";
constexpr char kDeviceTPU[] = "TPU";
constexpr char kTPUReplicatedCore[] = "TPU_REPLICATED_CORE";
constexpr char kTPUReplicatedHost[] = "TPU_REPLICATED_HOST";
constexpr char kBadIntArrayElementMsg[] =
    "bad '{0}' attribute at index {1}, not an int";

// in the input module op.
class PostCalibrationComponent : public Component {
 public:
  // Name of the post-training quantization post-calibration step. Used for
  // debugging purposes.
  static constexpr absl::string_view kName = "quant_ptq_post_calibration";

  explicit PostCalibrationComponent(absl::Nonnull<MLIRContext*> ctx);

  absl::StatusOr<ModuleOp> Run(
      ModuleOp module_op,

using ::tensorflow::Session;
using ::tensorflow::Status;
using ::tensorflow::StatusOr;
using ::tensorflow::Tensor;
using ::tensorflow::Var;

namespace {

constexpr char kResourceNameArgAttr[] = "tf.resource_name";
constexpr char kSavedModelArgAttr[] = "tf_saved_model.bound_input";

LogicalResult LiftVariablesFromSession(
    ModuleOp module, Session* session,
    const SmallSet<StringRef, 4>& resource_names) {