Wrong broadcasting of size 1 vector distributions defined with dims

[ricardoV94]

And this particular situation of not broadcasting two 1d arrays of lenghts 1 and 10 feels strange, how does broadcasting differ between Aesara and numpy?

Somehow the dims is turning the output variable into something that is not (known) to be broadcastable:

import pymc as pm

with pm.Model(coords={"slopes":["x_1"]}) as m:

    b_dim = pm.Normal("b_dim", sigma=1000, dims=("slopes"))
    b_size = pm.Normal("b_size", sigma=1000, size=1)

print(b_dim.broadcastable, b_size.broadcastable)
# (False,) (True,)

Originally posted by @ricardoV94 in #5561 (comment)

[ricardoV94]

The issue might come from this line, which implies the size is not fixed as it is wrapped in a aesara.shared:

pymc/pymc/model.py

Line 1113 in 08b6b01

self._dim_lengths[name] = length or aesara.shared(len(values))

I assume this is done to allow resizing by changing dims?

We need to enforce the broadcastable pattern of the shared variable We must not use shared variables for initial dims of size 1 or give users the tools to chose for themselves.

[ricardoV94]

CC @michaelosthege

[michaelosthege]

Yes, the aesara.shared is to enable a later re-sizing.

It was also my first suspicion that it's because of aesara.shared, but I dprinted the resulting variables and I didn't see SharedVariable there.
But maybe my analysis was not thorough enough.

[ricardoV94]

I assume this is what is happening?

import aesara
import aesara.tensor as at
x = at.random.normal(size=aesara.shared(1))
print(x.type)   # TensorType(float64, (None,))

[ricardoV94]

The shared variable is there in the PyMC distribution, just not obvious because it prints the type and it has no name:

aesara.dprint(d_dim)

normal_rv{0, (0, 0), floatX, False}.1 [id A] 'b_dim'   
 |RandomStateSharedVariable(<RandomState(MT19937) at 0x7FBD51C9CC40>) [id B]
 |SpecifyShape [id C] ''   
 | |MakeVector{dtype='int64'} [id D] ''   
 | | |<TensorType(int64, ())> [id E]
 | |TensorConstant{(1,) of 1} [id F]
 |TensorConstant{11} [id G]
 |TensorConstant{0} [id H]
 |TensorConstant{1000.0} [id I]

[id E] is the shared variable, I think

[ricardoV94]

I think we have to disable resizing of dims of length 1. Aesara needs the distinction as soon as we create the graph.

The value variables would also need to be recreated if the broadcastable pattern of the RV were to change

[ricardoV94]

@michaelosthege any input on this one?

[michaelosthege]

@michaelosthege any input on this one?

Sorry, no I'm not fluent enough to spontaneously be helpful, and I don't have the time to dig deeper right now

[ricardoV94]

The issue is caused by using shared variables for sizes defined via dims, the question is if it's fine to not allow that for the cases where the original dims have a size of 1, as that affects the broadcasting pattern of the variables.

[ricardoV94]

In the slack we proposed something like:

• All dims defined from coords via model.set_coord (this is the case for pm.Model(coords=...) ) should be TensorConstant
• All non-specified dims defined from pm.MutableData should be ScalarSharedVariable
  • Users can specify shapes/broadcastable dimensions for the underlying shared variables, so in that case we can consider them fixed.
• Everything else is TensorConstant