Same code that worked in PyMC 5.6.x fails in PyMC 5.8.x with `Rewrite failure due to: local_pow_to_nested_squaring`

Questions v5
1

Hi,

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

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.

I wonder if there is a way to solve this? Thank you!

2

Try to update to the very latest version. If it persists it’s a bug

3

Thanks @ricardoV94! We tried it on 5.8.2 and it still persists. We’ll submit an issue to pytensor.

2 Likes