Enforce positive domain of sqr(sqrt(x))

Author: jessegrabowski

pytensor/tensor/rewriting/math.py

@@ -424,12 +424,10 @@ def local_sqrt_sqr(fgraph, node):
 
     # Case for sqr(sqrt(x)) -> x
     if isinstance(prev_op, ps.Sqr) and isinstance(node_op, ps.Sqrt):
-        new_out = x.owner.inputs[0]
+        x = x.owner.inputs[0]
         old_out = node.outputs[0]
+        new_out = switch(ge(x, 0), x, np.asarray(np.nan, old_out.dtype))
 
-        # 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]