// TODO: b/230572023 - Consider improving shape inference for While op instead
  // of dropping the attribute. This need not be correct for models not trained
  // on TPU.

  // Optimizes TF graph via cleanups, merges, rewrites, constant folding,
  // and edge case handling where possible.
  pm.addNestedPass<func::FuncOp>(TF::CreateDropWhileShapeInvariantPass());
  pm.addNestedPass<func::FuncOp>(
      tf_executor::CreateTFExecutorGraphPruningPass());

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

// Rewrites RecvTPUEmbeddingActivationsOp and SendTPUEmbeddingGradients ops to
// internal variants by introducing XlaRecvTPUEmbeddingDeduplicationData op.
struct RewriteTPUEmbeddingOps
    : public impl::RewriteTPUEmbeddingOpsPassBase<RewriteTPUEmbeddingOps> {
  void runOnOperation() override;
};

// Rewrites the given op to `OpT` op after adding the given operand at the end.

  LogicalResult matchAndRewrite(FakeQuantOp op,
                                PatternRewriter &rewriter) const override {
    // TODO: If this pattern comes up more frequently, consider adding core
    // support for failable rewrites.
    if (failableRewrite(op, rewriter)) {
      *hadFailure = true;
      return failure();
    }
    return success();
  }

 private:
  bool *hadFailure;

	"os"
	"path/filepath"
	"sort"
	"strings"

	"cmd/internal/telemetry"
)

var (
	fset     = token.NewFileSet()
	exitCode = 0
)

var allowedRewrites = flag.String("r", "",
	"restrict the rewrites to this comma-separated list")

var forceRewrites = flag.String("force", "",
	"force these fixes to run even if the code looks updated")

var allowed, force map[string]bool

var (

};

// Quantization rewrite pattern using DQ as the root op.
struct StableHloQuantization
    : public StableHloQuantizationBase<StableHloQuantization> {
  explicit StableHloQuantization(MLIRContext* ctx)
      : StableHloQuantizationBase<StableHloQuantization>(ctx) {}
};

// Quantization rewrite pattern using Q as the root op. This is for the

 See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
// Rewrites computations generated by the xla.compile() Python code into
// XlaLaunch nodes.
//
// xla.compile() does two main things:
// a) marks operators that make up an XLA computation with the attribute

// some generic types that are legal in MHLO. This pass legalizes TF types into
// types that are legal in MHLO. For example, TF::Qint8Type is converted to i8.
// Rewrites here should run before TF to MHLO op legalizations are run.

#include <memory>

#include "absl/log/log.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/Support/Casting.h"

    auto filter_transpose_op = rewriter.create<mlir::stablehlo::TransposeOp>(
        op.getLoc(), /*resultType0=*/new_filter_tensor_type, /*operand=*/filter,
        rewriter.getDenseI64ArrayAttr(kOihwToHwioPermutation));

    // [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]
    const auto new_dimension_nums = rewriter.getAttr<ConvDimensionNumbersAttr>(

// Returns a TF Constant tensor with the passed in values.
TF::ConstOp GetI64ConstantTensor(PatternRewriter &rewriter,
                                 ArrayRef<int64_t> values, Location location) {
  auto cst_attr = rewriter.getI64TensorAttr(values);
  return rewriter.create<TF::ConstOp>(location, cst_attr.getType(), cst_attr);
}

// Rewrites broadcast->reshape to a reshape->broadcast that reduces
// the rank of the input and output of the broadcast.

type devirtualization struct {
	pos    string
	callee string
}

const profFileName = "devirt.pprof"
const preProfFileName = "devirt.pprof.node_map"

// testPGODevirtualize tests that specific PGO devirtualize rewrites are performed.
func testPGODevirtualize(t *testing.T, dir string, want []devirtualization, pgoProfileName string) {
	testenv.MustHaveGoRun(t)
	t.Parallel()

	const pkg = "example.com/pgo/devirtualize"