Change tolerance used to decide whether a constant is one in rewrite functions

[lciti]

Description

The previous tolerance used within a rewrite to decide whether a constant is one (or minus one) is too large.
For example c - sigmoid(x) is rewritten as sigmoid(-x) even when $c=1 − p$ where p is 1 in 10000.
Many rewrites currently use np.isclose and np.allclose with the default tolerances (rtol=1e-05, atol=1e-08), which are unnecessarily large (and independent on the data type of the constant computed).

This PR implements a function isclose used within all rewrites in place of np.isclose and np.allclose. This new function uses a much smaller tolerance by default, i.e. 10 unit in the last place (ULPs). This tolerance is dtype dependent, so it's stricter for a float64 than a float32. See #1497 for a back of the envelope justification for choosing 10 ULPs.

This PR also implements allow_cast in PatternNodeRewriter to allow rewrites that would otherwise fail when the new and old dtype differ. For example, a rewrite attempt for np.array(1., "float64") - sigmoid(x) (where x is fmatrix) currently fails because in the rewrite sigmoid(-x) the type would change. This PR allows an automatic cast to be added so the expression is rewritten as cast(sigmoid(-x), "float64").

Relevant tests added.

Related Issue

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

Checklist

• Checked that the pre-commit linting/style checks pass
• Included tests that prove the fix is effective or that the new feature works
• Added necessary documentation (docstrings and/or example notebooks)
• If you are a pro: each commit corresponds to a relevant logical change

Type of change

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):

---

📚 Documentation preview 📚: https://pytensor--1526.org.readthedocs.build/en/1526/

[ricardoV94]

Not all types have a dtype, we should check it's a TensorType before even trying to access dtype and doing stuff with it. I would perhaps write like this:

The whole logic is weird though with the if ret.owner, why do we care about the type of outputs we're not replacing. It's actually dangerous to try to replace only one of them without the user consent. Since this is WIP I would change to if len(node.outputs) != 1: return False, before we try to unify.

Then here we just have to worry about the final else branch below:

[old_out] = node.outputs
if not old_out.type.is_super(ret.type):
  if not (
    self.allow_cast 
    and isinstance(old_out.type, TensorType) 
    and isinstance(ret.type, TensorType)
  ):
    return False

  # Try to cast
  ret = ret.astype(old_out.type.dtype)
  if not old_out.type.is_super(ret.type):
    return False

[lciti]

I am happy to replace as you suggest but I am not sure how to fit it within the rest. This is the current code:

        if ret.owner:
            if not (
                len(node.outputs) == len(ret.owner.outputs)
                and all(
                    o.type.is_super(new_o.type)
                    for o, new_o in zip(node.outputs, ret.owner.outputs, strict=True)
                )
            ):
                return False
        else:
            # ret is just an input variable
            assert len(node.outputs) == 1
            if not node.outputs[0].type.is_super(ret.type):
                return False

[ricardoV94]

you only need what I wrote, above, template something like this

def transform(...):

...

if node.op != self.op:
    return False

if len(node.outputs) != 1:
  # PatternNodeRewriter doesn't support replacing multi-output nodes  
  return False

...

if not self.allow_multiple_clients:
  ...


# New logic

[old_out] = node.outputs
if not old_out.type.is_super(ret.type):
  # Type doesn't match
  if not (
    self.allow_cast 
    and isinstance(old_out.type, TensorType) 
    and isinstance(ret.type, TensorType)
  ):
    return False

  # Try to cast tensors
  ret = ret.astype(old_out.type.dtype)
  if not old_out.type.is_super(ret.type):
    # Still doesn't match
    return False

return [ret]

[lciti]

Are you sure PatternNodeRewriter is supposed to only work with single inputs? I get the following error:

    def test_patternsub_different_output_lengths():
        # Test that PatternNodeRewriter won't replace nodes with different numbers of outputs
        ps = PatternNodeRewriter(
            (op1, "x"),
            ("x"),
            name="ps",
        )
        rewriter = in2out(ps)
    
        x = MyVariable("x")
        e1, e2 = op_multiple_outputs(x)
        o = op1(e1)
    
        fgraph = FunctionGraph(inputs=[x], outputs=[o])
        rewriter.rewrite(fgraph)
>       assert fgraph.outputs[0].owner.op == op1
E       assert OpMultipleOutputs == op1
E        +  where OpMultipleOutputs = OpMultipleOutputs(x).op
E        +    where OpMultipleOutputs(x) = OpMultipleOutputs.0.owner

[ricardoV94]

I don't think that test makes sense. It's like saying you don't want to replace log(exp(x), if x comes from a multi-output node. We usually don't care about the provenance of a root variable in a rewrite. Nothing in that rewrite cares about op_multiple_outputs

[ricardoV94]

It was here: https://github.com/aesara-devs/aesara/pull/803/files

The problem was before the zip would be shorter if node.outputs and replacement didn't match in length. But the whole thing goes away if you just say it doesn't support replacing multiple outputs nodes, which it doesn't really.

That test can be removed in favor of one where it refuses to replace OpMultipleOutputs

[lciti]

Thanks. It sorts of makes sense to me but I know too little of the PyTensor internals to fully understand.
Can you propose a quick way to modify/replace the test with one where it refuses to replace OpMultipleOutputs?

[ricardoV94]

If you push your changes (if you haven't already), I can push the new test on top of it

[lciti]

I have pushed all my changes.

[ricardoV94]

I pushed a commit that changes the behavior of the test, have a look and let me know if there's anything else missing

[codecov]

Codecov Report

Attention: Patch coverage is 88.23529% with 2 lines in your changes missing coverage. Please review.

Project coverage is 81.85%. Comparing base (0bb15f9) to head (17fbc90).
Report is 39 commits behind head on main.

Files with missing lines	Patch %	Lines
pytensor/graph/rewriting/basic.py	85.71%	0 Missing and 1 partial ⚠️
pytensor/tensor/rewriting/math.py	90.00%	0 Missing and 1 partial ⚠️

❌ Your patch check has failed because the patch coverage (88.23%) is below the target coverage (100.00%). You can increase the patch coverage or adjust the target coverage.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1526      +/-   ##
==========================================
- Coverage   81.99%   81.85%   -0.14%     
==========================================
  Files         231      230       -1     
  Lines       52253    52532     +279     
  Branches     9203     9345     +142     
==========================================
+ Hits        42843    42999     +156     
  Misses       7099     7099              
- Partials     2311     2434     +123     

Files with missing lines	Coverage Δ
pytensor/graph/rewriting/basic.py	69.38% <85.71%> (-0.13%)	⬇️
pytensor/tensor/rewriting/math.py	89.65% <90.00%> (+0.03%)	⬆️

... and 33 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.