BUG: sampling fails with batched MvGaussianRandomWalk distributions #7075
Description
Describe the issue:
Sampling using MvGaussianRandomWalk with batch dimensions (i.e. 3-dimensional shape) fails with what appears to be a broadcasting issue. The distribution can be instantiated and evaluated, but not used to estimate a model. Details below.
Reproduceable code example:
import pymc as pm
import numpy as np
with pm.Model() as m:
delta = pm.MvGaussianRandomWalk("delta", mu=np.zeros(10), cov=np.eye(10), shape=(5, 7, 10))
trace = pm.sample()Error message:
ERROR:pytensor.graph.rewriting.basic:Rewrite failure due to: constant_folding
ERROR:pytensor.graph.rewriting.basic:node: Assert{msg=Could not broadcast dimensions. Broadcasting is only allowed along axes that have a statically known length 1. Use `specify_broadcastable` to inform PyTensor of a known shape.}(6, False)
ERROR:pytensor.graph.rewriting.basic:TRACEBACK:
ERROR:pytensor.graph.rewriting.basic:Traceback (most recent call last):
File "/usr/local/lib/python3.10/dist-packages/pytensor/graph/rewriting/basic.py", line 1922, in process_node
replacements = node_rewriter.transform(fgraph, node)
File "/usr/local/lib/python3.10/dist-packages/pytensor/graph/rewriting/basic.py", line 1082, in transform
return self.fn(fgraph, node)
File "/usr/local/lib/python3.10/dist-packages/pytensor/tensor/rewriting/basic.py", line 1106, in constant_folding
required = thunk()
File "/usr/local/lib/python3.10/dist-packages/pytensor/graph/op.py", line 518, in rval
r = p(n, [x[0] for x in i], o)
File "/usr/local/lib/python3.10/dist-packages/pytensor/raise_op.py", line 99, in perform
raise self.exc_type(self.msg)
AssertionError: Could not broadcast dimensions. Broadcasting is only allowed along axes that have a statically known length 1. Use `specify_broadcastable` to inform PyTensor of a known shape.PyMC version information:
PyMC 3.10.2, PyTensor 2.18.3
Context for the issue:
No response