mul(a,a) -> square(a)

sonos · Apr 12, 2024 · eebb56e · eebb56e
1 parent 60ad0c8
commit eebb56e
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/core/src/ops/math/mod.rs b/core/src/ops/math/mod.rs
@@ -331,6 +331,9 @@ fn declutter_mul(
     model: &TypedModel,
     node: &TypedNode,
 ) -> TractResult<Option<TypedModelPatch>> {
+    if node.inputs[0] == node.inputs[1] {
+        return Ok(Some(TypedModelPatch::replace_single_op(model, node, &node.inputs[0..1], square())?))
+    }
     if let Some(p) = declutter_neutral(model, node, 1, true).context("decluttering neutral")? {
         return Ok(Some(p));
     }