std::unique_ptr<OperationPass<func::FuncOp>> CreateTransformEinsumPass();

// Optimizes Tensorflow graph.
std::unique_ptr<OperationPass<func::FuncOp>> CreateTFOptimizePass();
void RegisterTFOptimizePassPipeline();

// Creates pass to rewrite RecvTPUEmbeddingActivationsOp and
// SendTPUEmbeddingGradients ops to internal variants.

}

def TensorFlowOptimizePass : Pass<"tf-optimize", "mlir::func::FuncOp"> {
  let summary = "Optimize TensorFlow module";

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

def PromoteResourcesToArgsPass : Pass<"tf-promote-resources-to-args", "ModuleOp"> {
  let summary = "Promote resources reads/writes to function inputs/outputs.";
  let description = [{