module->emitError(
        absl::StrCat("failed to fetch device manager: ", status.message()));
    return failure();
  }

  SmallVector<func::FuncOp, 2> session_init_funcs =
      tf_saved_model::GetInitializerFunctions(module);
  func::FuncOp session_init_func =
      session_init_funcs.empty() ? nullptr : session_init_funcs[0];

  TF::ResourceAnalyzer analyzer(module, /*skip_session_init=*/true);

      total_element_count += input_type.getNumElements();
    }
  }
  return total_element_count;
}

struct GetOpCostPass
    : mlir::PassWrapper<GetOpCostPass, OperationPass<func::FuncOp>> {
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(GetOpCostPass)

  llvm::StringRef getArgument() const final { return "tfl-get-op-cost"; }
  llvm::StringRef getDescription() const final {

// as it provides access to the TAC module.
template <typename T>
class TacFunctionPass : public TacPass<func::FuncOp> {
 public:
  using TacPass<func::FuncOp>::TacPass;

  ~TacFunctionPass() override = default;

  mlir::func::FuncOp getFunction() { return getOperation(); }

  virtual void runOnFunction() = 0;

  void runOnOperation() final {

    }
  }

  // TODO(bfontain): Check that the region here is the region
  // of the loop body func.
  // Find the FuncOp for the surrounding while loop body.
  func::FuncOp loop_body_func =
      (*forward_pass_ops.begin())->getParentOfType<func::FuncOp>();

  // merged_set will keep track of which ops are to be avoided when gather ops
  // for inclusion into the four extracted functions.

    * A tensor is quantized using a `func::FuncOp` whose name contains
      "uniform_quantize". The first argument is the tensor to be quantized,
      the second argument is the zero point constant (element type: int) and
      the third argument is the inverse scale constant (element type: float).
    * A tensor is dequantized using a `func::FuncOp` whose name contains

}

def TestResourceAliasAnalysis : Pass<"tf-test-resource-alias-analysis", "mlir::func::FuncOp"> {
  let summary = "Add remarks based on resource alias analysis result, for "
    "testing purpose.";
  let constructor = "::mlir::tf_test::CreateTestResourceAliasAnalysisPass()";
}

def TestClusteringPolicyPass : Pass<"tf-test-clustering-policy", "mlir::func::FuncOp"> {
  let summary = "Test pass for clustering based on the user defined policy";

    ModuleOp module,
    llvm::SmallDenseMap<StringRef, FuncToRewrite>& funcs_to_rewrite) {
  // Find functions containing communication ops.
  SmallVector<func::FuncOp, 4> funcs_to_visit;
  for (func::FuncOp func : module.getOps<func::FuncOp>()) {
    FuncToRewrite func_to_rewrite{/*original=*/func, /*control_flow_ops=*/{},
                                  /*control_flow_blocks=*/{},
                                  /*clone=*/nullptr};

include "mlir/Pass/PassBase.td"

def QuantConvertConst : Pass<"quant-convert-const", "func::FuncOp"> {
  let summary = "Converts constants followed by qbarrier to actual quantized "
                "values";
  let constructor = "mlir::quantfork::createConvertConstPass()";
}

def QuantConvertSimulatedQuant
    : Pass<"quant-convert-simulated-quantization", "func::FuncOp"> {
  let summary = "Converts training-time simulated quantization ops to "

class GroupByDialectPass
    : public impl::GroupByDialectPassBase<GroupByDialectPass> {
 public:
  void runOnOperation() override;

 private:
  void processFunction(mlir::func::FuncOp func, int& counter,
                       llvm::SmallDenseSet<StringRef>& dialects,
                       Operation* module);
  void processRegion(mlir::Region& region, int& counter,

// whose value is broacast to all replicas.
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateHoistBroadcastReadPass();

// Creates a pass that moves broadcasts from TF host ops to XLA code, encoded as
// XlaAllReduces. This enables use of the device network for broadcasts, which
// is faster.
std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
CreateXlaBroadcastPass();