From eebb56ec81ad4e0a7e7c0eb00c2b892327a56f4e Mon Sep 17 00:00:00 2001
From: Mathieu Poumeyrol <kali@zoy.org>
Date: Fri, 12 Apr 2024 09:15:02 +0200
Subject: [PATCH] mul(a,a) -> square(a)

---
 core/src/ops/math/mod.rs | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/core/src/ops/math/mod.rs b/core/src/ops/math/mod.rs
index 80559dc704..3c12de1886 100644
--- 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));
     }