PyMC3 elementwise if condition

Questions
1

I have the following model, where I have a feature space X of shape (NxTxK) which should be multiplied elementwise by alpha or beta depending on the sign of the value in phi which is also (NxTxK). The parameters alpha and beta are drawn from a prior distribution.
With numpy, the syntax would be: X = np.where(phi == 1, X*alpha, X*beta).
So far I used the following code which however gave me an error message.

with pm.Model() as model:

    alpha = pm.Beta(name="alpha", alpha=0.5, beta=0.5)
    beta = pm.Beta(name="beta", alpha=0.5, beta=0.5)

    X= pm.Deterministic(name="X", var=pm.math.where(pm.math.eq(phi, 1), X* alpha, X * beta))

    b= pm.Normal(name="b", mu=0, sigma=1, shape=n_vars_1)
  
    eps = pm.InverseGamma(name="eps", alpha=9.0, beta=4.0)

    y_hat = pm.Deterministic(name="y_hat", var=pm.math.sum(X * b, axis=2)
    y_like = pm.Normal(name="y_like", mu=y_hat, sd=eps, observed=y)
2

What is the error message you are getting?

3

TypeError: Condition given to the op has to be a scalar with 0 standing for False, anything else for True

4

Try aesara.tensor.switch

5

where is a provided alias to switch so it should have worked. Just to make sure check with theano.tensor.switch or aesara.tensor.switch depending on which version you are using.

What is phi?

1 Like
6

Phi is a TensorVariable that is drawn from a Bernoulli.
If I determine phi before initializing the model as a numpy array the code works.
Could it be that where or switchdoes not work for TensorVariables?

7

Those should definitely work with TensorVariables. Can you show the complete code, including the creation of Phi?