Fix bug in logprob_dimshuffle

[ricardoV94]

The logic to unshuffle the dimensions was wrong, but none of the pre-existing tests would reveal it.

The new test condition fails without the fix in multivariate test, and would also have failed in the univariate one, if local_dimshuffle_rv_lift was disabled.

Is there a way to disable the local_dimshuffle_rv_lift rewrite for the purposes of the test?

[brandonwillard]

Is there a way to disable the local_dimshuffle_rv_lift rewrite for the purposes of the test?

Yeah, we just need to create and use a Mode object with the desired rewrites (e.g. Mode(optimizer=...)). From there, Mode.excluding can be used to remove specific named rewrites.

[ricardoV94]

Is there a way to disable the local_dimshuffle_rv_lift rewrite for the purposes of the test?

Yeah, we just need to create and use a Mode object with the desired rewrites (e.g. Mode(optimizer=...)). From there, Mode.excluding can be used to remove specific named rewrites.

But right now there's no way of passing this to factorized_joint_logprob right?

[brandonwillard]

But right now there's no way of passing this to factorized_joint_logprob right?

Ah, no, but you can add some keyword options to construct_ir_fgraph for that.

[ricardoV94]

But right now there's no way of passing this to factorized_joint_logprob right?

Ah, no, but you can add some keyword options to construct_ir_fgraph for that.

Added an option

[ricardoV94]

This looks like a new Aesara-related bug: https://github.com/aesara-devs/aeppl/actions/runs/3067496963/jobs/4953853364

Opened PR to fix it: aesara-devs/aesara#1190