Add set_outfeed_config in XLA HloInstruction.

PiperOrigin-RevId: 317222410
Change-Id: I5de8a5067f1002a9d656d4e26d145ffe3fe372ed
This commit is contained in:
A. Unique TensorFlower 2020-06-18 18:15:12 -07:00 committed by TensorFlower Gardener
parent 18ca7c95d4
commit 723751b20e
3 changed files with 8 additions and 0 deletions

View File

@ -3908,6 +3908,10 @@ const string& HloInstruction::outfeed_config() const {
return Cast<HloOutfeedInstruction>(this)->outfeed_config();
}
void HloInstruction::set_outfeed_config(const string& config) {
return Cast<HloOutfeedInstruction>(this)->set_outfeed_config(config);
}
const std::vector<ReplicaGroup>& HloInstruction::replica_groups() const {
return Cast<HloCollectiveInstruction>(this)->replica_groups();
}

View File

@ -1755,6 +1755,9 @@ class HloInstruction {
// Returns the config for the Outfeed instruction.
const string& outfeed_config() const;
// Delegates to HloOutfeedInstruction::set_outfeed_config.
void set_outfeed_config(const string& config);
// Returns the shape for the Outfeed instruction.
const Shape& outfeed_shape() const;

View File

@ -1141,6 +1141,7 @@ class HloOutfeedInstruction : public HloInstruction {
const Shape& outfeed_shape() const { return outfeed_shape_; }
// Returns the config for the Outfeed instruction.
const string& outfeed_config() const { return outfeed_config_; }
void set_outfeed_config(const string& config) { outfeed_config_ = config; }
// Returns a serialized representation of this instruction.
HloInstructionProto ToProto() const override;