constexpr char kXLAShardingAttr[] = "_XlaSharding";
constexpr char kShardingAttr[] = "sharding";

typedef std::unordered_map<std::string, TF::TPUReplicateMetadataOp> MetadataMap;

struct TPUValidateInputsPass
    : public impl::TPUValidateInputsPassBase<TPUValidateInputsPass> {
  void runOnOperation() override;
};
bool IsTpuRegularOp(Operation* op) {
  static auto* ops = [] {
    llvm::SmallDenseSet<mlir::TypeID, 32>* ops_set =

  let summary = "Rewrite all TPU Partitioned ops into their V2 counterparts.";

  let constructor = "TFTPU::CreateTPUPartitionedOpConversionPass()";
}

def TPUValidateInputsPass : Pass<"tf-tpu-validate-inputs", "ModuleOp"> {
  let summary = "Validates inputs to the TPU TF/XLA bridge";

  let description = [{
    This pass checks that the IR has valid input to TPU TF/XLA bridge.