Rewrite sqr(sqrt(x)) -> |x| and sqrt(sqr(x)) -> x

Author: jessegrabowski

pytensor/tensor/rewriting/math.py

@@ -400,6 +400,39 @@ def local_exp_log(fgraph, node):
         return [exp(x)]
 
 
+@register_canonicalize
+@register_specialize
+@node_rewriter([sqrt, sqr])
+def local_sqrt_sqr(fgraph, node):
+    x = node.inputs[0]
+
+    if not (x.owner and isinstance(x.owner.op, Elemwise)):
+        return
+
+    prev_op = x.owner.op.scalar_op
+    node_op = node.op.scalar_op
+
+    # Case for sqrt(sqr(x)) -> |x|
+    if isinstance(prev_op, ps.Sqrt) and isinstance(node_op, ps.Sqr):
+        new_out = pt_abs(x.owner.inputs[0])
+        old_out = node.outputs[0]
+
+        # Handle potential integer to float cast by sqr
+        if new_out.dtype != old_out.dtype:
+            new_out = cast(new_out, old_out.dtype)
+        return [new_out]
+
+    # Case for sqr(sqrt(x)) -> x
+    if isinstance(prev_op, ps.Sqr) and isinstance(node_op, ps.Sqrt):
+        new_out = x.owner.inputs[0]
+        old_out = node.outputs[0]
+
+        # Handle potential integer to float cast by sqrt
+        if x.dtype != old_out.dtype:
+            new_out = cast(new_out, old_out.dtype)
+        return [new_out]
+
+
 @register_specialize
 @node_rewriter([exp, expm1])
 def local_exp_log_nan_switch(fgraph, node):

tests/tensor/rewriting/test_math.py

@@ -2031,6 +2031,56 @@ def test_exp_log_nested(self, nested_expression, expected_switches):
         assert len(ops_graph) == expected_switches
 
 
+class TestSqrSqrt:
+    def setup_method(self):
+        mode = get_default_mode()
+        self.mode = mode.including(
+            "local_sqrt_sqr",
+        ).excluding("fusion")
+        self.rng = np.random.default_rng()
+
+    def test_sqr_sqrt(self):
+        # sqrt(x) ** 2 -> x
+        data = self.rng.normal(size=(4, 3)).astype(config.floatX)
+        x = pt.tensor("x", shape=(None, None))
+        f = function([x], sqr(sqrt(x)), mode=self.mode)
+        graph = f.maker.fgraph.toposort()
+        assert not any(
+            isinstance(node.op, Elemwise)
+            and isinstance(node.op.scalar_op, ps.Sqr | ps.Sqrt)
+            for node in graph
+        )
+        assert any(
+            isinstance(node.op, Elemwise) and isinstance(node.op.scalar_op, ps.Abs)
+            for node in graph
+        )
+
+        np.testing.assert_array_equal(f(data), np.abs(data))
+
+    def test_sqrt_sqr(self):
+        data = self.rng.normal(size=(4, 3)).astype(config.floatX)
+        x = pt.tensor("x", shape=(None, None))
+        f = function([x], sqrt(sqr(x)), mode=self.mode)
+        graph = f.maker.fgraph.toposort()
+        assert not any(
+            isinstance(node.op, Elemwise)
+            and isinstance(node.op.scalar_op, ps.Sqr | ps.Sqrt)
+            for node in graph
+        )
+        np.testing.assert_array_equal(f(data), data)
+
+    def test_sqr_sqrt_integer_upcast(self):
+        x = ivector("x")
+        f = function([x], sqrt(sqr(x)), mode=self.mode)
+        nodes = f.maker.fgraph.toposort()
+        assert not any(
+            isinstance(node.op, Elemwise)
+            and isinstance(node.op.scalar_op, ps.Sqr | ps.Sqrt)
+            for node in nodes
+        )
+        assert f([1, 2, 3]).dtype in ["float32", "float64"]
+
+
 class TestLocalSwitchSink:
     def setup_method(self):
         # condition values