diff --git a/tensorflow/compiler/xla/service/hlo_instructions.cc b/tensorflow/compiler/xla/service/hlo_instructions.cc index bcc00d806da..2a53841fd34 100644 --- a/tensorflow/compiler/xla/service/hlo_instructions.cc +++ b/tensorflow/compiler/xla/service/hlo_instructions.cc @@ -550,6 +550,7 @@ bool HloCollectiveInstruction::IdenticalSlowPath( const auto& casted_other = static_cast<const HloCollectiveInstruction&>(other); return HloChannelInstruction::IdenticalSlowPath(other, eq_computations) && + constrain_layout() == casted_other.constrain_layout() && absl::c_equal(replica_groups(), casted_other.replica_groups(), [](const ReplicaGroup& a, const ReplicaGroup& b) { return absl::c_equal(a.replica_ids(), b.replica_ids()); @@ -1101,7 +1102,9 @@ bool HloMapInstruction::IdenticalSlowPath( const HloInstruction& other, const std::function<bool(const HloComputation*, const HloComputation*)>& eq_computations) const { - return eq_computations(to_apply(), other.to_apply()); + const auto& casted_other = static_cast<const HloMapInstruction&>(other); + return eq_computations(to_apply(), casted_other.to_apply()) && + dimensions() == casted_other.dimensions(); } std::unique_ptr<HloInstruction> HloMapInstruction::CloneWithNewOperandsImpl( @@ -2515,7 +2518,8 @@ bool HloDynamicSliceInstruction::IdenticalSlowPath( const HloInstruction& other, const std::function<bool(const HloComputation*, const HloComputation*)>& eq_computations) const { - return true; + const auto& casted_other = static_cast<const HloMapInstruction&>(other); + return dynamic_slice_sizes() == casted_other.dynamic_slice_sizes(); } std::unique_ptr<HloInstruction>