BUG: sampling error when indexing over mutliple dimensions

[flo-schu]

Describe the issue:

When indexing/broadcasting a random variable over multiple dimensions, during sampling, a dimension mismatch is called. However, when evaluating the final array, it matches the dimensionality of the input array. The error does not occur, when indexing is only done over one dimension. I'm not sure where the problem lies. Or even preferably, whether there is a better solution for the task.

Reproduceable code example:

import pymc as pm
import numpy as np
rng = np.random.default_rng(5648)

# prior dimensions
N = 1
S = 2
T = 3

# indexer
a = np.repeat(np.arange(S), [1,2])
b = np.repeat(np.arange(T), [1,2,1])
K = len(a)
L = len(b)

y = rng.standard_normal((N,K,L))

with pm.Model() as model:
    mu = pm.Normal("mu", size=(N, S, T))

    mua = mu[:, a, :]
    mub = mua[:, :, b]

    y = pm.Normal("y", mub, sigma=1, observed=y)

    prior = pm.sample_prior_predictive()

    y.eval()
    pm.sample()

Error message:

pymc.sampling.parallel.RemoteTraceback: 
"""
Traceback (most recent call last):
  File "/home/schunckf/.conda/envs/pymc/lib/python3.10/site-packages/aesara/compile/function/types.py", line 971, in __call__
    self.vm()
AssertionError: SpecifyShape: dim 1 of input has shape 3, expected 1.

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/schunckf/.conda/envs/pymc/lib/python3.10/site-packages/pymc/sampling/parallel.py", line 123, in run
    self._start_loop()
  File "/home/schunckf/.conda/envs/pymc/lib/python3.10/site-packages/pymc/sampling/parallel.py", line 176, in _start_loop
    point, stats = self._step_method.step(self._point)
  File "/home/schunckf/.conda/envs/pymc/lib/python3.10/site-packages/pymc/step_methods/arraystep.py", line 273, in step
    return super().step(point)
  File "/home/schunckf/.conda/envs/pymc/lib/python3.10/site-packages/pymc/step_methods/arraystep.py", line 199, in step
    apoint, stats = self.astep(q)
  File "/home/schunckf/.conda/envs/pymc/lib/python3.10/site-packages/pymc/step_methods/hmc/base_hmc.py", line 168, in astep
    start = self.integrator.compute_state(q0, p0)
  File "/home/schunckf/.conda/envs/pymc/lib/python3.10/site-packages/pymc/step_methods/hmc/integration.py", line 56, in compute_state
    logp, dlogp = self._logp_dlogp_func(q)
  File "/home/schunckf/.conda/envs/pymc/lib/python3.10/site-packages/pymc/model.py", line 409, in __call__
    cost, *grads = self._aesara_function(*grad_vars)
  File "/home/schunckf/.conda/envs/pymc/lib/python3.10/site-packages/aesara/compile/function/types.py", line 984, in __call__
    raise_with_op(
  File "/home/schunckf/.conda/envs/pymc/lib/python3.10/site-packages/aesara/link/utils.py", line 534, in raise_with_op
    raise exc_value.with_traceback(exc_trace)
  File "/home/schunckf/.conda/envs/pymc/lib/python3.10/site-packages/aesara/compile/function/types.py", line 971, in __call__
    self.vm()
AssertionError: SpecifyShape: dim 1 of input has shape 3, expected 1.
Apply node that caused the error: SpecifyShape(Elemwise{sub,no_inplace}.0, NoneConst, TensorConstant{1}, NoneConst)
Toposort index: 13
Inputs types: [TensorType(float64, (None, 3, 4)), <aesara.tensor.type_other.NoneTypeT object at 0x2afa257583a0>, TensorType(int8, ()), <aesara.tensor.type_other.NoneTypeT object at 0x2afa257583a0>]
Inputs shapes: [(1, 3, 4), 'No shapes', (), 'No shapes']
Inputs strides: [(96, 32, 8), 'No strides', (), 'No strides']
Inputs values: ['not shown', None, array(1, dtype=int8), None]
Outputs clients: [[InplaceDimShuffle{2,1,0}(SpecifyShape.0)]]

HINT: Re-running with most Aesara optimizations disabled could provide a back-trace showing when this node was created. This can be done by setting the Aesara flag 'optimizer=fast_compile'. If that does not work, Aesara optimizations can be disabled with 'optimizer=None'.
HINT: Use the Aesara flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.
"""

The above exception was the direct cause of the following exception:

AssertionError: SpecifyShape: dim 1 of input has shape 3, expected 1.
Apply node that caused the error: SpecifyShape(Elemwise{sub,no_inplace}.0, NoneConst, TensorConstant{1}, NoneConst)
Toposort index: 13
Inputs types: [TensorType(float64, (None, 3, 4)), <aesara.tensor.type_other.NoneTypeT object at 0x2afa257583a0>, TensorType(int8, ()), <aesara.tensor.type_other.NoneTypeT object at 0x2afa257583a0>]
Inputs shapes: [(1, 3, 4), 'No shapes', (), 'No shapes']
Inputs strides: [(96, 32, 8), 'No strides', (), 'No strides']
Inputs values: ['not shown', None, array(1, dtype=int8), None]
Outputs clients: [[InplaceDimShuffle{2,1,0}(SpecifyShape.0)]]

HINT: Re-running with most Aesara optimizations disabled could provide a back-trace showing when this node was created. This can be done by setting the Aesara flag 'optimizer=fast_compile'. If that does not work, Aesara optimizations can be disabled with 'optimizer=None'.
HINT: Use the Aesara flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "<string>", line 29, in <module>
  File "/home/schunckf/.conda/envs/pymc/lib/python3.10/site-packages/pymc/sampling/mcmc.py", line 529, in sample
    mtrace = _mp_sample(**sample_args, **parallel_args)
  File "/home/schunckf/.conda/envs/pymc/lib/python3.10/site-packages/pymc/sampling/mcmc.py", line 1019, in _mp_sample
    for draw in sampler:
  File "/home/schunckf/.conda/envs/pymc/lib/python3.10/site-packages/pymc/sampling/parallel.py", line 447, in __iter__
    draw = ProcessAdapter.recv_draw(self._active)
  File "/home/schunckf/.conda/envs/pymc/lib/python3.10/site-packages/pymc/sampling/parallel.py", line 332, in recv_draw
    raise error from old_error
pymc.sampling.parallel.ParallelSamplingError: Chain 0 failed with: SpecifyShape: dim 1 of input has shape 3, expected 1.
Apply node that caused the error: SpecifyShape(Elemwise{sub,no_inplace}.0, NoneConst, TensorConstant{1}, NoneConst)
Toposort index: 13
Inputs types: [TensorType(float64, (None, 3, 4)), <aesara.tensor.type_other.NoneTypeT object at 0x2afa257583a0>, TensorType(int8, ()), <aesara.tensor.type_other.NoneTypeT object at 0x2afa257583a0>]
Inputs shapes: [(1, 3, 4), 'No shapes', (), 'No shapes']
Inputs strides: [(96, 32, 8), 'No strides', (), 'No strides']
Inputs values: ['not shown', None, array(1, dtype=int8), None]
Outputs clients: [[InplaceDimShuffle{2,1,0}(SpecifyShape.0)]]

HINT: Re-running with most Aesara optimizations disabled could provide a back-trace showing when this node was created. This can be done by setting the Aesara flag 'optimizer=fast_compile'. If that does not work, Aesara optimizations can be disabled with 'optimizer=None'.
HINT: Use the Aesara flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.

PyMC version information:

4.4

Context for the issue:

No response

[ricardoV94]

Thanks for the clean reproducible bug report!

Here is a more direct way to trigger the error

model.compile_logp()(model.initial_point())  # AssertionError: SpecifyShape: Got shape (1, 3, 4), expected (3, 3, 4).

Tracking issue: pymc-devs/pytensor#98

[ricardoV94]

For now I think you can just remove the rightmost semicolons in mua:

import pymc as pm
import numpy as np

rng = np.random.default_rng(5648)

# prior dimensions
N = 1
S = 2
T = 3

# indexer
a = np.repeat(np.arange(S), [1,2])
b = np.repeat(np.arange(T), [1,2,1])
K = len(a)
L = len(b)

y = rng.standard_normal((N,K,L))

with pm.Model() as model:
    mu = pm.Normal("mu", size=(N, S, T))
    mua = mu[:, a]
    mub = mua[:, :, b]
    y = pm.Normal("y", mub, sigma=1, observed=y)

model.compile_logp()(model.initial_point())

[flo-schu]

This solved problem, also in a more complex application! Lovely :)

[ricardoV94]

Thanks for the bug report. Closing this as it should be solved by pymc-devs/pytensor#101