What is the proper way to do einsum in Pymc3?

Questions version agnostic
1

I am trying to convert the following code (written in tensorflow) to Pymc3 which is from the Book; Bayesian Modeling and Computation in Python, Martin, Kumar & Lao, chapter 6, listing 6.3.

    y_hat = (intercept[..., None] +
             tf.einsum("ij,...->...i", trend, trend_coeff) +
             tf.einsum("ij,...j->...i", seasonality, seasonality_coeff))

I am especially having trouble writing tf.einsum into Pymc3 inside the pymc model …
However, the full code block is as follows

tfd = tfp.distributions
root = tfd.JointDistributionCoroutine.Root

@tfd.JointDistributionCoroutine
def ts_regression_model():
    intercept = yield root(tfd.Normal(0., 100., name="intercept"))
    trend_coeff = yield root(tfd.Normal(0., 10., name="trend_coeff"))

    seasonality_coeff = yield root(
        tfd.Sample(tfd.Normal(0., 1.),
                   sample_shape=seasonality.shape[-1],
                   name="seasonality_coeff"))
    noise = yield root(tfd.HalfCauchy(loc=0., scale=5., name="noise_sigma"))
    y_hat = (intercept[..., None] +
             tf.einsum("ij,...->...i", trend, trend_coeff) +
             tf.einsum("ij,...j->...i", seasonality, seasonality_coeff))
    observed = yield tfd.Independent(
        tfd.Normal(y_hat, noise[..., None]),
        reinterpreted_batch_ndims=1,
        name="observed")

Also, I have heard Theano’s tensor_dot can help; but is there a another way around?

2

I will let @aloctavodia @RavinKumar or @junpenglao chime in, but in the meantime, this is a pretty thorough discussion of einsum.

1 Like
3

aesara/pymc does not have einsum implemented yet, you can use tensordot instead (might need to transpose the input)

1 Like
4

Ok after standardizing the outputs and features I did something like this; but the posteriors were wrong!!

with pm.Model() as model_1:
  intercept = pm.Normal("intercept",0,0.2)
  trend_coeff=pm.Normal('trend_coeff',0.0,0.5)
  seasonality_coeff=pm.Normal('seasonality_coeff',0.0,0.5,shape=seasonality.shape[-1])
  noise=pm.Exponential('noise',1)
  y_hat=intercept[...,None] +trend@trend_coeff +tensor.tensordot(seasonality_coeff,seasonality,(0,1))
  y= pm.Normal('y',mu=y_hat,sigma=noise,observed=co2_by_month_training_data.CO2.values)
  trace_1=pm.sample(cores=2)

Is there a problem with how I specified the model?