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]

I added that test, so let me see if I can dig up the rationale

[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]

If you are not evaluating f, just rewrite it with rewrite_graph, possibly including ("canonicalize", "stabilize", "specialize"), or whatever is needed

[ricardoV94]

You can do pytensor.dprint(tuple[Variable]). If you want the rewritten, expected, which many times I do while writing these sort of tests we could add an assert_equal_computations helper that does that. That way it's reusable and doesn't make each test very verbose like this?

[lciti]

I agree assert_equal_computations might be useful (would that go into tests/unittest_tools.py ?). I think it would be useful to know the original expression too, what do yo think?

[lciti]

Something like:

def assert_equal_computations(
    rewritten,
    expected,
    *args,
    original=None,
    **kwargs
):
    """
    Assert that `rewritten` computes the same as `expected`.
    
    Parameters
    ----------
    rewritten
        The expression after the rewrite pass.
    expected
        The reference expression to compare against.
    *args, **kwargs
        Extra arguments forwarded to equal_computations.
    original : optional
        If given, will be printed in the error message.
    """
    ok = equal_computations(rewritten, expected, *args, **kwargs)

    if not ok:
        parts = []

        def _dprint(expr):
            return pytensor.dprint(expr, print_type=True, file="str")

        if original is not None:
            parts.append(f"Original:\n{_dprint(original)}")
        parts.append(f"Rewritten:\n{_dprint(rewritten)}")
        parts.append(f"Expected:\n{_dprint(expected)}")

        raise AssertionError("\n\n".join(parts))

    return True

[ricardoV94]

Looks lovely

[ricardoV94]

I was reviewing commit by commit, see you changed this after. Anyway my original comment still stands.

Generally, feel free to squash commits and force-push when iterating on a PR so the git changes stay clean