[Grappler] Fix a bug in folding multiply into add arithmetic optimization
PiperOrigin-RevId: 308348068 Change-Id: I61ca637344519f9effee313ea3f911c462d9bfe6
This commit is contained in:
parent
687a2d4bc5
commit
3b1545a555
@ -2178,6 +2178,8 @@ class FoldMultiplyIntoConv : public ArithmeticOptimizerStage {
|
|||||||
// Check that value preserving chain is the only consumer of the Mul output.
|
// Check that value preserving chain is the only consumer of the Mul output.
|
||||||
TF_RETURN_IF_TRUE(!IsAnyMul(*source));
|
TF_RETURN_IF_TRUE(!IsAnyMul(*source));
|
||||||
TF_RETURN_IF_TRUE(NumNonControlOutputs(*source, *ctx().node_map) != 1);
|
TF_RETURN_IF_TRUE(NumNonControlOutputs(*source, *ctx().node_map) != 1);
|
||||||
|
// And that Mul is not in the preserve set.
|
||||||
|
TF_RETURN_IF_TRUE(IsInPreserveSet(*source));
|
||||||
|
|
||||||
const NodeDef* mul = source;
|
const NodeDef* mul = source;
|
||||||
int input_idx = 0;
|
int input_idx = 0;
|
||||||
|
Loading…
Reference in New Issue
Block a user