BUG: sampling fails with batched MvGaussianRandomWalk distributions

[fonnesbeck]

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