auto input_ty = mlir::cast<ShapedType>(input.getType());
    auto min_ty = mlir::cast<ShapedType>(min.getType());
    auto max_ty = mlir::cast<ShapedType>(max.getType());

    if (!input_ty.hasRank() || !min_ty.hasRank() || !max_ty.hasRank()) {
      return failure();
    }

    auto shape = rewriter.create<TF::ShapeOp>(
        op.getLoc(),
        tensorflow::GetTypeFromTFTensorShape({input_ty.getRank()},