using mlir::Type;
using mlir::Value;
using mlir::ValueRange;
using mlir::WalkResult;
using mlir::func::FuncOp;
using mlir::TF::ConstOp;
using mlir::TF::FillOp;
using mlir::TF::IdentityOp;
using mlir::TF::XlaAllReduceOp;
using mlir::tf_device::ClusterOp;
using mlir::tf_device::LaunchOp;
using mlir::tf_device::ParallelExecuteOp;
using mlir::tf_device::ReplicateOp;

#define GEN_PASS_DEF_XLABROADCASTPASS

};

// Converts XlaAllReduce. Not thread-safe.
class ConvertXlaAllReduce
    : public CollectiveRewritePattern<TF::XlaAllReduceOp> {
 public:
  using CollectiveRewritePattern::CollectiveRewritePattern;

  LogicalResult matchAndRewrite(TF::XlaAllReduceOp all_reduce,
                                PatternRewriter& rewriter) const override {
    DenseIntElementsAttr replica_groups;

    TypeID::get<TF::UnsortedSegmentProdOp>(),
    TypeID::get<TF::UnsortedSegmentSumOp>(),
    TypeID::get<TF::XdivyOp>(),
    TypeID::get<TF::XlaSendTPUEmbeddingGradientsOp>(),
    TypeID::get<TF::XlaAllReduceOp>(),
    TypeID::get<TF::XlaGatherOp>(),
    TypeID::get<TF::Xlog1pyOp>(),
    TypeID::get<TF::XlogyOp>(),
    TypeID::get<TF::ZerosLikeOp>(),
    TypeID::get<TF::ZetaOp>(),
  };
  // clang-format on