A little help with broadcasting

I have two distributions with different shapes (eq and region). I thought I could use broadcasting to combine them together to compute n. However, I couldn’t reshape them to create something that aesara thought was broadcastable. Eventually, I used the code below and it seems to have worked. Is there a better way to do this?

mag = pm.Normal(
        "mag",
        mu=events.loc[id_eq, "mag"].values,
        sigma=events.loc[id_eq, "mag_sd"].values,
        dims="eq",
    )
n_a = pm.LogNormal("n_a", mu=np.log(0.6884), sigma=0.05, dims="region")

foo = n_b * pm.math.tanh(mag - n_c)
n = at.stack([n_a[0] * foo, n_a[1] * foo])

What is the shape of your desired output? Scalar operations like multiplication are not defined between two vectors of different shapes. Calling the shape of n_a as (n, 1) and the shape of foo as (m, 1), it looks like you are taking the outer product so that n has shape (n, m), then raveling it into a single (n * m, 1) vector.

You could do this using broadcasting by adding extra dimensions to your variables, i.e. n_a[:, None] * foo[None, :] will have shape (n, m). You can also use at.outer to get the same result without having to think about the indexing.

1 Like

The intended shape would be [region, eq] in size. Thanks! Both of those solutions work. I had previously tried to use .reshape(1, n_region) to change the shape, which I thought was consistent with [None, :].

Reshape will add a dimension in Aesara, but only if you give the shape as a tuple. So .reshape(1, -1) doesn’t work, but .reshape((1, -1)) does. This is a rare difference between aesara and numpy syntax.

1 Like