void runOnOperation() override;
};

// Returns true iff the broadcast val can be substituted with an XlaAllReduce.
// Sets `zero` and `shape` as params for the dummy zeros that will be created.
bool GetDummyParams(OpBuilder& builder, Value val_bcast, Attribute& zero,
                    DenseIntElementsAttr& shape) {
  Type type = val_bcast.getType();
  Type elem_type = getElementTypeOrSelf(type);