BUG: Rewrite failure local_pow_to_nested_squaring when static type shape changes

[digicosmos86]

Describe the issue:

We are using PyMC to build a model with a custom distribution and likelihood for the hssm package. The code ran fine with 5.6.x, but after updating PyMC dependency to 5.8.x, the code broke under some circumstances with certain parameter settings at the sampling stage with this error:

ERROR (pytensor.graph.rewriting.basic): Rewrite failure due to: local_pow_to_nested_squaring
ERROR (pytensor.graph.rewriting.basic): node: Pow(True_div.0, [3])
ERROR (pytensor.graph.rewriting.basic): TRACEBACK:
ERROR (pytensor.graph.rewriting.basic): Traceback (most recent call last):
  File "/Users/yxu150/HSSM/.venv/lib/python3.9/site-packages/pytensor/graph/rewriting/basic.py", line 1922, in process_node
    replacements = node_rewriter.transform(fgraph, node)
  File "/Users/yxu150/HSSM/.venv/lib/python3.9/site-packages/pytensor/graph/rewriting/basic.py", line 1082, in transform
    return self.fn(fgraph, node)
  File "/Users/yxu150/HSSM/.venv/lib/python3.9/site-packages/pytensor/tensor/rewriting/math.py", line 2139, in local_pow_to_nested_squaring
    assert rval[0].type == node.outputs[0].type, (rval, node.outputs)
AssertionError: ([Composite{(sqr(i0) * i0)}.0], [Pow.0])

Reproducable code example:

Since there is only one line with pt.pow(*, 3), we found the offending code seems to be this line:

p = p / pt.sqrt(2 * np.pi * pt.power(tt, 3))

where tt can take negative values, which we thought might be where the composite type came from. However, even after ensuring that tt is only positive with something like tt = pt.maximum(tt, 1e-25), the problem persists. After changing optimizer setting to o2, the problem also goes away.

To reproduce the error, first, install hssm via:

pip install git+https://github.com/lnccbrown/HSSM.git@280-pin-numpy-version

The code that produces the error:

import hssm
import pymc as pm
import pytensor
from hssm.likelihoods import DDM

pytensor.config.floatX = "float32"

v_true, a_true, z_true, t_true = [0.5, 1.5, 0.5, 0.5]
dataset = hssm.simulate_data(
    model="ddm",
    theta=[v_true, a_true, z_true, t_true],
    size=1000,
).values

with pm.Model() as ddm_pymc:
    v = pm.Uniform("v", lower=-10.0, upper=10.0)
    a = pm.HalfNormal("a", sigma=2.0)
    z = pm.Uniform("z", lower=0.01, upper=0.99)
    t = pm.Uniform("t", lower=0.0, upper=0.6, initval=0.1)

    ddm = DDM("ddm", v=v, a=a, z=z, t=t, observed=dataset)

    ddm_pymc_trace = pm.sample()

Error message:

No response

PyTensor version information:

pymc = 5.8.2 pytensor = 2.16.2

Context for the issue:

No response

[ricardoV94]

You guys have quite a number of dependencies...

[ricardoV94]

This one is a bit tricky, the rewrite gets to the point where it tries to replace Pow with shape=(None,) by Sqr with shape=(1,):

Pow [id A] <Vector(float32, shape=(?,))>
 ├─ Mul [id B] <Vector(float32, shape=(1,))>
 │  ├─ [-1.] [id C] <Vector(float32, shape=(1,))>
 │  └─ ExpandDims{axis=0} [id D] <Vector(float32, shape=(1,))>
 │     └─ Add [id E] <Scalar(float32, shape=())>
 └─ [2] [id F] <Vector(int8, shape=(1,))>

Sqr [id G] <Vector(float32, shape=(1,))>
 └─ Mul [id B] <Vector(float32, shape=(1,))>
    └─ ···

The fix is not to fail, but to get out graciously.

For the time being, do you have any idea where that surprising parameter with shape=(1,) is coming from? I'm calling it surprising because PyTensor did not know that was the shape of Pow. To properly fix this and still apply the rewrite, we need something like #408 so as to not change the meaning of the graph after the rewriting.

If you know you are creating stuff with shape=(1,) in advance and that PyTensor doesn't know this, you can wrap those objects in x = pt.specify_shape=(x, shape=(1,)). That way the rewrite won't complain by the time it gets there.

[ricardoV94]

Reproducible snippet:

import pytensor
import pytensor.tensor as pt
from pytensor.graph.basic import Apply

x = pt.vector("x", shape=(1,))
# Create a node manually that does not know the output should have `shape=(1,)`
node = Apply(pt.pow, [x, pt.as_tensor([2.0])], [pt.tensor(shape=(None,))])
y = node.outputs[0]

y.eval({x: 1.0})  # AssertionError

[ricardoV94]

Should be patched by #461

[digicosmos86]

Thanks @ricardoV94! We don't know where that surprising parameter comes from. We'll take a look at what is responsible for this