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!!