#include "tensorflow/core/util/device_name_utils.h"

namespace tensorflow {
using tsl::StatusOr;

inline constexpr absl::string_view kNumCoresPerReplicaAttr =
    "num_cores_per_replica";
inline constexpr absl::string_view kTopologyAttr = "topology";
inline constexpr absl::string_view kDeviceAssignmentAttr = "device_assignment";

// A TPU device for execution alongside its associated host CPU device.

#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<

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";

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";

inline constexpr char kDebugModeOpQuantAttrName[] = "debug_quant";

// Used to annotate custom ops if they are quantizable.
inline constexpr char kQuantTraitAttrName[] = "_tfl_quant_trait";
enum QuantizationTrait { FullyQuantizable = 0, NotQuantizable = 1 };
inline constexpr absl::string_view QuantTraitValues[] = {"fully_quantizable",
                                                         "not_quantizable"};

inline constexpr double kNearZeroTolerance = 1.0e-6;

namespace tensorflow {
namespace internal {
namespace {

using ::tensorflow::quantization::PyFunctionLibrary;
using ::tflite::optimize::ReducedPrecisionSupport;

// Op def string for TFLite_Detection_PostProcess Op.
constexpr mlir::StringRef kDetectionPostProcessOp =
    "name: 'TFLite_Detection_PostProcess' input_arg: { name: "
    "'raw_outputs/box_encodings' type: DT_FLOAT } input_arg: { name: "

        AddUse(some_operation(rng), some_tensor(rng));
      }
    }
  };
};

TEST_F(RematSimulationTest, SimulationAgreesWithReality) {
  constexpr int kNumOperations = 128;
  constexpr int kNumTensors = 32;
  constexpr int kNumUses = kNumOperations * kNumTensors / 4;

  std::mt19937 rng;
  for (int i = 0; i < 1024; ++i) {
    RandomRemat remat(kNumOperations, kNumTensors, kNumUses, rng);

using ::tsl::protobuf::TextFormat;

// Default version number for native serialization.
constexpr int64_t kDefaultVersion = 9;
// Default platform for XlaCallModuleOp.
constexpr StringRef kPlatformCpu = "CPU";
// Name of `tf.XlaCallModule`'s dictionary attribute for keeping the
// deserialized stablehlo module's attributes.
constexpr StringRef kStablehloModuleAttrsAttrName = "_stablehlo_module_attrs";

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

namespace mlir {
namespace TFTPU {

namespace {

constexpr char kDeviceAttr[] = "device";
constexpr char kFuncDeviceAttr[] = "tf.device";
constexpr char kDefaultShardingValue[] = "";
constexpr char kMirroredVariableIndicesAttr[] = "_mirrored_variable_indices";

std::string GetRandomStateVariableName() {