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

class DropWhileShapeInvariantPass
    : public impl::DropWhileShapeInvariantPassBase<
          DropWhileShapeInvariantPass> {
  void runOnOperation() override;
};

#define GEN_PASS_DEF_DROPWHILESHAPEINVARIANTINDEVICECLUSTERPASS

    is in training mode or in evaluation mode.
  }];

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

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

  let description = [{