BUG: the tolerance used to decide whether a constant is one (or minus one) in rewrite functions may be too large

[lciti]

Describe the issue:

I think that the tolerance used to decide whether a constant is one (or minus one) is too large.
We can argue whether this is a bug or not but I think it may lead to issues.
See code below.

The issue is the use of np.isclose and np.allclose within functions _is_1 in pytensor/tensor/rewriting/math.py, is_1pexp, is_neg`, etc.

The default tolerances used are rtol=1e-05, atol=1e-08, which are unnecessarily large (and independent on the data type of the constant computed).

I think it's not unreasonable to have a constant $(1-p)$ where $p$ is the probability of a relatively rare event (e.g. 1 in 10000) and the last thing we want is for that constant to be treated as 1.

I propose:

• using a much smaller tolerance by default
• make this tolerance dependent on the data type.
• implementing an in-house version of isclose/allclose to be used where needed (so it can be reimplemented by the user and its effects apply to all simplifications)

Le's take for example, the following expression:
$\big(\cos(7)^2 + \sin(7)^2\big) \Big(\frac12 + \frac13 + \frac16 \Big) \big((\sqrt 2)^2 - 1\big)^{\pi}$

I did some quick rounding‐error analysis of this expression and the error bound should be approximately 30 deltas. Assuming they are roughly independent, this means roughly 6 eps. In practice, I have tried it and I get an error of 5 eps in float64 and 3 eps in float32. So I think it's safe to take something like 10 eps as tolerance when assessing if something is one (or minus one), which corresponds to 1e-6 for float32 and 2e-15 for float64.

Reproducable code example:

import pytensor
import pytensor.tensor as pt
from pytensor.graph import rewrite_graph

x = pt.scalar("x")
out = 0.99999 - pt.sigmoid(x)

pytensor.dprint(out, print_type=True)
print('='*20)
with pytensor.config.change_flags(optimizer_verbose = True):
    fn = pytensor.function([x], out, mode="FAST_RUN")
    #fn = rewrite_graph(out, include=("canonicalize", "stabilize", "specialize"))

print('='*20)
pytensor.dprint(fn, print_type=True);

Error message:

Sub [id A] <Scalar(float64, shape=())>
 ├─ 0.99999 [id B] <Scalar(float64, shape=())>
 └─ Sigmoid [id C] <Scalar(float64, shape=())>
    └─ x [id D] <Scalar(float64, shape=())>
====================
rewriting: rewrite local_1msigmoid replaces Sub.0 of Sub(0.99999, Sigmoid.0) with Sigmoid.0 of Sigmoid(Neg.0)
rewriting: rewrite FusionOptimizer replaces Sigmoid.0 of Sigmoid(Neg.0) with Composite{sigmoid((-i0))}.0 of Composite{sigmoid((-i0))}(x)
====================
Composite{sigmoid((-i0))} [id A] <Scalar(float64, shape=())> 0
 └─ x [id B] <Scalar(float64, shape=())>

Inner graphs:

Composite{sigmoid((-i0))} [id A]
 ← sigmoid [id C] <float64> 'o0'
    └─ neg [id D] <float64>
       └─ i0 [id E] <float64>

PyTensor version information:

2.31.3

Context for the issue:

No response

[ricardoV94]

Yes we can be more strict

[lciti]

@ricardoV94 I was looking into this and I have a question.
A rewrite like "c - sigmoid(x) → sigmoid(-x)" where c is a float64 constant (scalar or matrix) close or equal to one and x is a float32 variable (scalar or matrix) is not performed. After some digging I found that this is probably due to o.type.is_super(new_o.type) that prevents a rewrite from changing the type of the variable being computed. This makes sense. However, I find it a bit strange that in the following:

x = fmatrix()
out1 = 1 - pt.sigmoid(x)
out2 = 1. - pt.sigmoid(x)
out3 = np.array(1.) - pt.sigmoid(x)

the last expression (out3) ends up being much less accurate than the first two. Should we do anything about this? (Nothing / address it / warn the user of the failed rewrite when optimizer_verbose == True)

[ricardoV94]

We could add an option to the PatternNodeRewriter to allow upcasting the output pattern to the original type so that the rewrite applies.

The casting operation by itself can be expensive (a whole array iteration and copy), but if it is happening alongside another Elemwise operation it ends up being basically free because of loop fusion.

Regarding informing users. In general when working with float32 you may want to set pytensor.config.floatX = "float32" and the flag pytensor.config.warn_float64 = "warn", although that is designed for people who want to avoid float64, it could also be useful in this scenario to make sure you are not accidentally mixing dtypes.

import pytensor.tensor as pt
pytensor.config.floatX = "float32"
pytensor.config.warn_float64 = "warn"
x = pt.fmatrix()
out = np.array(1.) - pt.sigmoid(x)

Otherwise, I use optimizer_verbose a lot and I'm afraid this may pollute the output. We could add more fine-grained control though.

[lciti]

Thanks for your reply.
I agree that upcasting may be the best option. It would be virtually free thanks to it being fused with other elemwise operations and it's a cost that was already there I presume (i.e. without the rewrtie the float32 output of sigmoid would go through a cvtss2sd before bing added to the double constant 1.0).
This is a bit beyond my ability I am afraid.

For the optimizer_verbose, are you sure it would pollute significantly? It could only kick in for a stabilize rewrite and only when everything matches except for the type.

[lciti]

I have given it a try but I am not sure if it breaks anything else. I was thinking something along the lines of:

        if ret.owner:
            if len(node.outputs) != len(ret.owner.outputs):
                return
            if len(node.outputs) > 1:
                if not all(
                        o.type.is_super(new_o.type) for o, new_o in zip(
                            node.outputs, ret.owner.outputs, strict=True)
                ):
                    return False
            if ret.owner.outputs[0].type.dtype != node.outputs[0].type.dtype and self.allow_upcast:
                ret = cast(ret, node.outputs[0].type.dtype)
            if not node.outputs[0].type.is_super(ret.owner.outputs[0].type):
                return False

(I am only upcasting in the case of a single output because I have no idea how to do it for multiple outputs).
Here self.allow_upcast (or self.auto_cast or similar) would be a new argument to PatternNodeRewriter.