TF2 branch doesn't work

Hi I hope this is the right place to post about other dev branches (couldn't find an Issue page there).

Was going to experiment with the TF2 branch njtierney:tf2-poke-tf-fun. To update from greta v0.4.2, I did

devtools::install_github("njtierney/greta", rep = "tf2-poke-tf-fun")

but greta_sitrep returns messages saying that the old TF1.x stuff are still loading in the environment (sorry I lost the messages now).

So I ran install_greta_deps(), which then led to greta_sitrep returning

✔ python (v3.7) available     
✔ TensorFlow (v2.9.1) available   
✔ TensorFlow Probability (v0.17.0) available  
✔ greta conda environment available            
✔ Initialising python and checking dependencies ... done!

But when I ran the examples, I got the following error:

> x <- rnorm(10)
> mu <- normal(0, 5)
> sigma <- lognormal(1, 0.1)
> distribution(x) <- normal(mu, sigma)
> m <- model(mu, sigma)
> 
> # carry out mcmc on the model
> draws <- mcmc(m, n_samples = 100)
Error in py_call_impl(callable, dots$args, dots$keywords) : 
  RuntimeError: Evaluation error: AttributeError: module 'tensorflow_probability.python.bijectors' has no attribute 'AffineScalar'
.

By the way, I tried this out to see if it indirectly resolves the issue related to initials() here #440 . When I tried to sample from a model containing lkj_correlation() arrays, it now returns a different kind of error:

Error in py_call_impl(callable, dots$args, dots$keywords) : 
  RuntimeError: Evaluation error: ValueError: can't split axis of size 2344 into pieces of size [1,4,1,4,2156,6,4,1,4,128,6,4,1,4] for '{{node split}} = SplitV[T=DT_DOUBLE, Tlen=DT_INT32, num_split=14](free_state, Const, split/split_dim)' with input shapes: [1,2344], [14], [] and with computed input tensors: input[1] = <1 4 1 4 2156 6 4 1 4 128 6 4 1 4>, input[2] = <1>.
.

Hi @hrlai

Thanks for posting, and for trying out the TF2 branch!

Yes, you'll need to install the TF 2 dependencies again when you change to the TF2 branch, we will make the process hopefully friendly, and loudly help people when we get the TF2 code working.

Yes at the moment the branch, https://github.com/greta-dev/greta/pull/534, is still experimental. We were able to get MCMC to run, but it was very slow and had poor convergence. greta was previously written around the graph-based structure of TF1, and now that TF2 uses eager evaluation, we need to change a lot of the internals of greta to support that, and also ensure that we appropriately construct the graphs so they don't get recreated each time. There's still a fair bit of work that we need to do, but I'll keep you posted as we get further along. For the time being I'll close this issue as we are tracking the TF2 progress over in https://github.com/greta-dev/greta/issues/306 and https://github.com/greta-dev/greta/issues/458

greta-dev / greta

TF2 branch doesn't work #539