// in this pass. Not thread-safe.
};

// 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 {