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

class DropWhileShapeInvariantInDeviceClusterPass
    : public impl::DropWhileShapeInvariantInDeviceClusterPassBase<
          DropWhileShapeInvariantInDeviceClusterPass> {
  void runOnOperation() override;
};

void DropWhileShapeInvariantAttr(Operation* op) {
  if (llvm::isa<WhileOp, WhileRegionOp>(op))

    these ops. This is only safe to do when compiling to XLA.
  }];

  let constructor = "TF::CreateDropWhileShapeInvariantPass()";
}

def DropWhileShapeInvariantInDeviceClusterPass : Pass<"tf-drop-while-shape-invariant-in-device-cluster", "mlir::func::FuncOp"> {
  let summary = "Drop `shape_invariant` attribute from While/WhileRegion ops inside device cluster.";

  let description = [{