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

struct ReplicateInvariantOpHoistingPass
    : public impl::ReplicateInvariantOpHoistingPassBase<
          ReplicateInvariantOpHoistingPass> {
  void runOnOperation() override;
};

// Check if op directly uses a key in `virtual_devices`.

  }];

  // Required for mhlo bounded shape extension.
  let dependentDialects = ["mhlo::MhloDialect"];
  let constructor = "TFTPU::CreateTPUAnnotateDynamicShapeInputsPass()";
}

def ReplicateInvariantOpHoistingPass : Pass<"tf-replicate-invariant-op-hoisting", "mlir::func::FuncOp"> {
  let summary = "Hoists replicate invariant operations out of replicate";

  let description = [{