Code conversion Bayesian Modeling and Computation in Python ch6. Help on AR(1) model implementation

Questions
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.14gam_with_latent_ar. for better understanding.

def generate_gam_ar_latent(training=True):

    @tfd.JointDistributionCoroutine
    def gam_with_latent_ar():
        seasonality, trend, noise_sigma = yield from gam_trend_seasonality()
        
        # Latent AR(1)
        ar_sigma = yield root(tfd.HalfNormal(.1, name="ar_sigma"))
        rho = yield root(tfd.Uniform(-1., 1., name="rho"))
        def ar_fun(y):
            loc = tf.concat([tf.zeros_like(y[..., :1]), y[..., :-1]],
                            axis=-1) * rho[..., None]
            return tfd.Independent(
                tfd.Normal(loc=loc, scale=ar_sigma[..., None]),
                reinterpreted_batch_ndims=1)
        temporal_error = yield tfd.Autoregressive(
            distribution_fn=ar_fun,
            sample0=tf.zeros_like(trend),
            num_steps=trend.shape[-1],
            name="temporal_error")

        # Linear prediction
        y_hat = seasonality + trend + temporal_error
        if training:
            y_hat = y_hat[..., :co2_by_month_training_data.shape[0]]

        # Likelihood
        observed = yield tfd.Independent(
            tfd.Normal(y_hat, noise_sigma[..., None]),
            reinterpreted_batch_ndims=1,
            name="observed"
        )

    return gam_with_latent_ar

gam_with_latent_ar = generate_gam_ar_latent()

I have already converted the trend and seasonality and it seems to be working fine but I am stuck with the AR part.

P.S thank you @RavinKumar @junpenglao, and @aloctavodia for writing this awesome book and making it accessible online.

Also if anyone else is working on converting chapter 6 of the Book for better understanding; please let me know!!

2

Did you try the pm.AR distribution in v4? It should work out of the box

2 Likes
3

Thank you @junpenglao. Any idea how if not when Pymc4 will start to work in google colab?

4

I think I should make a different Topic for the availability of Pymc4 in google colab!!

5

A few Google folks are working on it (me, @RavinKumar and @colcarroll)

1 Like
6

It would be great!!