PyMC3 append tensor variables

I try to store 4 2D random variables of shape (98, 26) in a container with desired shape (98, 26, 4) within a pm.Model() environment.

In numpy:

state = np.random.randint(low=0, high=4, size=(98, 26, 4))
B = np.random.random(size=(4, 4))
B_s = np.empty(shape=(98, 26, 4))
for i in range(0, 4):
    B_s[:, :, i] = B[:, i][state[:, :, i]]

However, this does not work as the objects are theano tensors.

The model looks as follows:

with pm.Model() as mdoel:

 B = pm.Normal(name="B", mu=0, sigma=1, shape=(4, 4))
 B_s = np.empty(shape=(98, 26, 4))
 for i in range(0, 4):
        B_s[:, :, i] = B[:, i][state[:, :, i]]

y_hat = np.sum(X * B_s, axis=2)

where X ist of shape (98, 26, 4).

The error is
ValueError: setting an array element with a sequence.

Hi there,

I think the way to go here is to use indices instead. Array indexing takes a bit of getting used to, but you can almost always avoid for loops and assignments. Here’s what seems to work for your case:

# Your code
state = np.random.randint(low=0, high=4, size=(98, 26, 4))
B = np.random.random(size=(4, 4))
B_s = np.empty(shape=(98, 26, 4))
for i in range(0, 4):
    B_s[:, :, i] = B[:, i][state[:, :, i]]

# Indexing alternative:
B_s_alt = B[state, np.arange(4)]

# Make sure this worked:
# Should return True
np.allclose(B_s_alt, B_s)

Then you can implement your model as follows:

with pm.Model() as mdoel:

 B = pm.Normal(name="B", mu=0, sigma=1, shape=(4, 4))
 B_s = B[state, np.arange(4)]

y_hat = np.sum(X * B_s, axis=2)

Hope that helps – let me know!

2 Likes