Tensorflow likelihood working with `tfp.vi.surrogate_posterior`

Output from this script with USE_LIBSBN = False:
2021-08-24 10:35:39.331431: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /usr/local/lib
2021-08-24 10:35:39.331450: I tensorflow/stream_executor/cuda/cudart_stub.cc:29] Ignore above cudart dlerror if you do not have a GPU set up on your machine.
2021-08-24 10:35:40.494823: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcuda.so.1'; dlerror: libcuda.so.1: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /usr/local/lib
2021-08-24 10:35:40.494843: W tensorflow/stream_executor/cuda/cuda_driver.cc:269] failed call to cuInit: UNKNOWN ERROR (303)
2021-08-24 10:35:40.494856: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:156] kernel driver does not appear to be running on this host (CS424460): /proc/driver/nvidia/version does not exist
2021-08-24 10:35:40.495076: I tensorflow/core/platform/cpu_feature_guard.cc:151] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
/home/cswa648/dev/treeflow/treeflow/coalescent.py:99: UserWarning: Dummy sampling
  warnings.warn("Dummy sampling")
Trials  [------------------------------------]    0%
Traceback (most recent call last):
  File "run.py", line 117, in <module>
    log_posterior, q_tmp, tf.optimizers.Adam(learning_rate=0.0001), 5
  File "/home/cswa648/miniconda3/envs/libsbn/lib/python3.7/site-packages/tensorflow_probability/python/vi/optimization.py", line 309, in fit_surrogate_posterior
    name=name)
  File "/home/cswa648/miniconda3/envs/libsbn/lib/python3.7/site-packages/tensorflow_probability/python/math/minimize.py", line 347, in minimize
    seed=seed)
  File "/home/cswa648/miniconda3/envs/libsbn/lib/python3.7/site-packages/tensorflow/python/util/traceback_utils.py", line 153, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/home/cswa648/miniconda3/envs/libsbn/lib/python3.7/site-packages/tensorflow_probability/python/math/minimize.py", line 88, in optimizer_step
    grads = tape.gradient(loss, watched_variables)
TypeError: _capture_helper() takes 3 positional arguments but 4 were given
christiaanjs / treeflow

Tensorflow likelihood working with `tfp.vi.surrogate_posterior` #36
