MCMC crashes using future::plan(cluster) or future::plan(multisession)

Hi,
Thanks a lot for this amazing package!

I'm trying to use greta's mcmc on a Linux (Ubuntu) machine in RStudio 2022.01.0 Build 205 and R 4.1.3.
To check if everything works I've copied the example in https://mdscheuerell.github.io/gretaDFA/

When I try the following command:

future::plan(cluster) # or future::plan(multisession)
mod_fit <- mcmc(mod_defn, verbose = TRUE,n_cores=8,
sampler = hmc(Lmin = 1, Lmax = 30, epsilon = 0.001, diag_sd = 1),
warmup = 2000, n_samples = 5000, thin = 10, chains = 1,
initial_values = initials(RR_est = runif(1, 0.1, 1),
sigma_est = runif(1, 0.1, 1)
)
)

I get the error

Loaded Tensorflow version 1.14.0
Error in py_call_impl(callable, dots$args, dots$keywords) :
RuntimeError: Evaluation error: object 'tf' not found.

Detailed traceback:
File "/home/gianni/.local/share/r-miniconda/envs/greta-env/lib/python3.7/site-packages/tensorflow_probability/python/mcmc/sample.py", line 326, in sample_chain
previous_kernel_results = kernel.bootstrap_results(current_state)

Any suggestion on how to fix this problem?
Thanks
Gianni

----------------------------------------
Reproducible chunk
```{r}
library(greta)
library(tensorflow)
library(MASS)
library(viridis)

pkg <- reticulate::import("pkg_resources")
pkg$get_distribution("tensorflow_probability")$version

# We need to define a helper function to fill in the diagonal of the prior for the loadings matrix Z with the appropriate distribution (see below).

zd <- function(M, sigma) {
zd <- zeros(M)
for(i in 1:M) {
zd[i] <- ld(i, M = M, sigma = sigma, dim = 1)
}
return(zd)
}
NN <- 30
TT <- 30
MM <- 3
## simulation
# latent factors

set.seed(123)
## MM x TT matrix of innovations
ww <- matrix(rnorm(MM*TT, 0, 1), MM, TT)
ww[,1] <- rnorm(MM, 0, sqrt(5))
## MM x TT matrix of scaled latent trends
xx <- t(scale(apply(ww,1,cumsum)))
head(xx)
dim(xx)
matplot(1:dim(xx)[2],t(xx),'l')

# loading matrix
ZZ <- matrix(runif(NN*MM, -1, 1), NN, MM)
diag(ZZ) <- rev(sort(abs(diag(ZZ))))
ZZ[upper.tri(ZZ)] <- 0
ZZ <- round(ZZ, 2)

# observed series
## obs var
obs_var <- 0.2^2
## obs errors
ee <- t(mvrnorm(TT, matrix(0,NN,1), diag(obs_var,NN,NN)))
## NN x TT matrix of observed data
yy <- ZZ %*% xx + ee
matplot(1:dim(yy)[2],t(yy),'l')

## ESTIMATION
source("~/Dropbox/Econometrics/R/Functions/LD_distribution/ld_distro.r")
## PRIORS
## empty loadings matrix
ZZ_est <- zeros(NN,MM)
## define sigma
sigma_est = normal(0, 2, truncation = c(0, Inf))
## diagonal
idx_d <- row(ZZ_est) == col(ZZ_est)
ZZ_est_raw_d = zd(MM, sigma_est)
ZZ_est[idx_d] <- ZZ_est_raw_d
## sub-diagonal
idx_s <- lower.tri(ZZ_est)
ZZ_est_raw_s = normal(0, sigma_est, dim = sum(idx_s), truncation = c(-1
,1))
ZZ_est[idx_s] <- ZZ_est_raw_s

# check that everything works
# head(calculate(ZZ_est, values=list(ZZ_est_raw_d = rep(1, sum(idx_d)),
# ZZ_est_raw_s = rep(2, sum(idx_s)))))

# COVARIANCE MATRIX
# diagonal
RR_est = inverse_gamma(alpha = 1, beta = 3, 1, truncation=c(0, 1))

## inital factor
X0 = normal(mean = 0, sd = sqrt(10), dim = c(MM, 1))
## Likerlihood function
## factors for t = 2-TT
XX = normal(mean = 0, sd = 1, dim = c(MM, TT))
## cumsum over proc errs for the random walk
xx_est <- t(apply(cbind(X0,XX), 1, "cumsum"))[,-1]

# prepare observed data
## scale data
yy_z <- t(scale(t(yy), scale = FALSE))
## vectorize data
yy_vec <- as_data(matrix(yy_z, 1, NN*TT))
## vectorize mean
Zx_vec <- t(c(ZZ_est %*% xx_est))
## define likelihood
distribution(yy_vec) = normal(Zx_vec, RR_est)

### FIT THE MODEL WITH MCMC
# define the model for greta
mod_defn <- model(xx_est, ZZ_est, RR_est, sigma_est)
#### mcmc
# require(future) # all options of plan crash (except sequential)
# cores=detectCores()-1
# cl1 <- makeCluster(cores)
# # registerDoParallel(cl1)
# plan(multisession) # the cluster plan conflicts with TensorFlow
mod_fit <- mcmc(mod_defn, verbose = TRUE,n_cores=8,
sampler = hmc(Lmin = 1, Lmax = 30, epsilon = 0.001, diag_sd = 1),
warmup = 2000, n_samples = 5000, thin = 10, chains = 1,
initial_values = initials(RR_est = runif(1, 0.1, 1),
sigma_est = runif(1, 0.1, 1)
# XX = matrix(rnorm(MM*TT), MM, TT),
# ZZ_est_raw_s = matrix(rnorm(sum(idx_s),
# 0,
# 0.1),
# ncol = 1)
)
)
```