build_factored_surrogate_posterior() got an unexpected keyword argument 'batch_shape'

Hello,

I'm getting into trouble when i'm trying to run a tensorflow probability example found here:
https://www.tensorflow.org/probability/api_docs/python/tfp/sts/impute_missing_values

My packages:

Package Version Location

absl-py 0.12.0
access 1.1.3
affine 2.3.0
aiobotocore 1.4.1
aiohttp 3.7.4.post0
aioitertools 0.8.0
albumentations 1.0.3
alembic 1.7.3
alibi-detect 0.7.2
allennlp 2.7.0
altair 4.1.0
annoy 1.17.0
ansiwrap 0.8.4
anyio 3.3.0
appdirs 1.4.4
argcomplete 1.12.3
argon2-cffi 20.1.0
arrow 1.1.1
arviz 0.11.2
asn1crypto 1.4.0
astunparse 1.6.3
async-generator 1.10
async-timeout 3.0.1
attrs 21.2.0
audioread 2.1.9
autobahn 21.3.1
autocfg 0.0.8
Automat 20.2.0
autopage 0.4.0
aws-requests-auth 0.4.3
Babel 2.9.1
backcall 0.2.0
backports.csv 1.0.7
backports.functools-lru-cache 1.6.4
backports.zoneinfo 0.2.1
base58 2.1.0
bayesian-optimization 1.2.0
bayespy 0.5.22
beatrix-jupyterlab 0.9.1
beautifulsoup4 4.10.0
binaryornot 0.4.4
biopython 1.79
black 21.7b0
bleach 4.0.0
blinker 1.4
blis 0.7.4
bokeh 2.3.3
Boruta 0.3
boto3 1.18.44
botocore 1.21.44
Bottleneck 1.3.2
bq-helper 0.4.1 /src/bq-helper
bqplot 0.12.30
branca 0.4.2
brewer2mpl 1.4.1
brotlipy 0.7.0
cachetools 4.2.2
caip-notebooks-serverextension 1.0.0
Cartopy 0.19.0.post1
catalogue 1.0.0
catalyst 21.8
catboost 0.26.1
category-encoders 2.2.2
certifi 2021.5.30
cesium 0.9.12
cffi 1.14.6
cftime 1.5.0
chardet 4.0.0
charset-normalizer 2.0.4
checklist 0.0.11
cheroot 8.5.2
CherryPy 18.6.1
cleverhans 4.0.0
click 8.0.1
click-plugins 1.1.1
cliff 3.9.0
cligj 0.7.2
cloud-tpu-client 0.10
cloudpickle 1.6.0
cmaes 0.8.2
cmd2 2.2.0
cmdstanpy 0.9.5
cmudict 1.0.2
colorama 0.4.4
colorcet 2.0.6
colorlog 6.4.1
colorlover 0.3.0
conda 4.10.3
conda-package-handling 1.7.3
configparser 5.0.2
confuse 1.5.0
constantly 15.1.0
contextily 1.1.0
contextlib2 21.6.0
convertdate 2.3.2
cookiecutter 1.7.0
coverage 5.5
cryptography 3.4.7
cudf 21.8.3
cufflinks 0.17.3
cuml 21.8.2
cupy 9.4.0
CVXcanon 0.1.2
cycler 0.10.0
cymem 2.0.5
cysignals 1.10.3
Cython 0.29.24
cytoolz 0.11.0
dask 2021.9.0
dask-cudf 21.8.3
dataclasses 0.8
datasets 1.12.1
datashader 0.13.0
datashape 0.5.2
datatable 1.0.0
deap 1.3.1
debugpy 1.4.1
decorator 5.0.9
deepdish 0.3.6
defusedxml 0.7.1
Delorean 1.0.0
deprecation 2.1.0
descartes 1.1.0
dill 0.3.4
dipy 1.4.1
distributed 2021.9.0
dlib 19.22.1
dm-tree 0.1.6
dnspython 2.1.0
docker 5.0.0
docker-pycreds 0.4.0
docutils 0.17.1
earthengine-api 0.1.283
easydev 0.12.0
easydict 1.9
easyocr 1.4.1
ecos 2.0.7.post1
eli5 0.11.0
emoji 1.5.0
en-core-web-lg 2.3.1
en-core-web-sm 2.3.1
entrypoints 0.3
ephem 4.0.0.2
esda 2.4.1
essentia 2.1b6.dev374
fairscale 0.4.0
fastai 2.2.7
fastavro 1.4.4
fastcore 1.3.26
fastprogress 1.0.0
fastrlock 0.6
fasttext 0.9.2
fbpca 1.0
fbprophet 0.7.1
feather-format 0.4.1
featuretools 0.27.1
feedparser 6.0.8
filelock 3.0.12
Fiona 1.8.20
fitter 1.4.0
flake8 3.9.2
flashtext 2.7
Flask 2.0.1
Flask-Cors 3.0.10
Flask-JSON 0.3.4
flatbuffers 1.12
folium 0.12.1
fsspec 2021.8.1
funcy 1.16
fury 0.7.1
future 0.18.2
fuzzywuzzy 0.18.0
gast 0.3.3
gatspy 0.3
gcsfs 2021.7.0
gensim 4.0.1
geographiclib 1.52
Geohash 1.0
geojson 2.5.0
geopandas 0.9.0
geoplot 0.4.4
geopy 2.2.0
geoviews 1.9.1
ggplot 0.11.5
giddy 2.3.3
gitdb 4.0.7
GitPython 3.1.18
gluoncv 0.10.4.post4
gluonnlp 0.10.0
google-api-core 1.31.1
google-api-python-client 1.8.0
google-auth 1.34.0
google-auth-httplib2 0.1.0
google-auth-oauthlib 0.4.5
google-cloud-aiplatform 0.6.0a1
google-cloud-appengine-logging 0.1.4
google-cloud-audit-log 0.1.0
google-cloud-automl 1.0.1
google-cloud-bigquery 2.2.0
google-cloud-bigtable 2.3.3
google-cloud-core 1.7.2
google-cloud-dataproc 2.5.0
google-cloud-datastore 2.1.6
google-cloud-firestore 2.2.0
google-cloud-kms 2.5.0
google-cloud-language 2.2.2
google-cloud-logging 2.6.0
google-cloud-monitoring 2.4.2
google-cloud-pubsub 1.7.0
google-cloud-scheduler 2.3.2
google-cloud-spanner 3.8.0
google-cloud-speech 2.7.0
google-cloud-storage 1.42.0
google-cloud-tasks 2.5.1
google-cloud-translate 3.3.2
google-cloud-videointelligence 2.3.2
google-cloud-vision 2.4.2
google-crc32c 1.1.2
google-pasta 0.2.0
google-resumable-media 1.3.3
googleapis-common-protos 1.53.0
gplearn 0.4.1
gpxpy 1.4.2
graphviz 0.8.4
greenlet 1.1.1
grpc-google-iam-v1 0.12.3
grpcio 1.32.0
grpcio-gcp 0.2.2
gym 0.20.0
h2o 3.34.0.1
h5py 2.10.0
haversine 2.5.1
HeapDict 1.0.1
hep-ml 0.7.0
hijri-converter 2.2.1
hmmlearn 0.2.6
holidays 0.11.2
holoviews 1.14.6
hpsklearn 0.1.0
html5lib 1.1
htmlmin 0.1.12
httplib2 0.19.1
httplib2shim 0.0.3
huggingface-hub 0.0.17
humanize 3.11.0
hunspell 0.5.5
husl 4.0.3
hydra-slayer 0.2.0
hyperlink 21.0.0
hyperopt 0.2.5
hypertools 0.7.0
ibis-framework 1.4.0
idna 2.10
imagecodecs 2021.8.26
ImageHash 4.2.1
imageio 2.9.0
imbalanced-learn 0.8.0
imgaug 0.4.0
implicit 0.4.8
importlib-metadata 3.4.0
importlib-resources 5.2.2
incremental 21.3.0
inequality 1.0.0
iniconfig 1.1.1
ipykernel 6.2.0
ipympl 0.7.0
ipython 7.26.0
ipython-genutils 0.2.0
ipython-sql 0.3.9
ipywidgets 7.6.3
iso-639 0.4.5
iso3166 2.0.1
isoweek 1.3.3
itsdangerous 2.0.1
Janome 0.4.1
jaraco.classes 3.2.1
jaraco.collections 3.4.0
jaraco.functools 3.3.0
jaraco.text 3.5.1
jax 0.2.19
jaxlib 0.1.70+cuda110
jedi 0.18.0
jieba 0.42.1
Jinja2 3.0.1
jinja2-time 0.2.0
jmespath 0.10.0
joblib 1.0.1
json5 0.9.5
jsonlines 1.2.0
jsonnet 0.17.0
jsonschema 3.2.0
jupyter 1.0.0
jupyter-client 6.1.12
jupyter-console 6.4.0
jupyter-core 4.7.1
jupyter-http-over-ws 0.0.8
jupyter-server 1.10.2
jupyter-server-mathjax 0.2.3
jupyterlab 1.2.21
jupyterlab-git 0.11.0
jupyterlab-pygments 0.1.2
jupyterlab-server 1.2.0
jupyterlab-widgets 1.0.0
kaggle 1.5.12
kaggle-environments 1.8.12
Keras-Preprocessing 1.1.2
keras-tuner 1.0.4
kiwisolver 1.3.1
kmapper 2.0.1
kmodes 0.11.0
korean-lunar-calendar 0.2.1
kornia 0.5.8
kt-legacy 1.0.4
kubernetes 18.20.0
langid 1.1.6
learntools 0.3.4
leven 1.0.4
libcst 0.3.20
libpysal 4.5.1
librosa 0.8.1
lightfm 1.16
lightgbm 3.2.1
lime 0.2.0.1
line-profiler 3.3.0
llvmlite 0.36.0
lmdb 1.2.1
lml 0.1.0
locket 0.2.1
LunarCalendar 0.0.9
lxml 4.6.3
Mako 1.1.5
mapclassify 2.4.3
marisa-trie 0.7.7
Markdown 3.3.4
markovify 0.9.3
MarkupSafe 2.0.1
matplotlib 3.4.3
matplotlib-inline 0.1.2
matplotlib-venn 0.11.6
matrixprofile 1.1.10
mccabe 0.6.1
memory-profiler 0.58.0
mercantile 1.2.1
mgwr 2.1.2
missingno 0.4.2
mistune 0.8.4
mizani 0.7.3
ml-metrics 0.1.4
mlcrate 0.2.0
mlens 0.2.3
mlxtend 0.19.0
mmh3 3.0.0
mne 0.23.4
mnist 0.2.2
mock 4.0.3
more-itertools 8.10.0
mpld3 0.5.5
mpmath 1.2.1
msgpack 1.0.2
msgpack-numpy 0.4.7.1
multidict 5.1.0
multimethod 1.4
multipledispatch 0.6.0
multiprocess 0.70.12.2
munch 2.5.0
murmurhash 1.0.5
mxnet-cu110 1.8.0.post0
mypy-extensions 0.4.3
nb-conda 2.2.1
nb-conda-kernels 2.3.1
nbclient 0.5.4
nbconvert 6.1.0
nbdime 3.1.0
nbformat 5.1.3
nest-asyncio 1.5.1
netCDF4 1.5.7
networkx 2.5
nibabel 3.2.1
nilearn 0.8.1
nltk 3.2.4
nnabla 1.21.0
nnabla-ext-cuda110 1.21.0
nose 1.3.7
notebook 6.4.3
notebook-executor 0.2
numba 0.53.1
numexpr 2.7.3
numpy 1.19.5
nvtx 0.2.3
oauth2client 4.1.3
oauthlib 3.1.1
odfpy 1.4.1
olefile 0.46
opencv-contrib-python 4.5.3.56
opencv-python 4.5.3.56
opencv-python-headless 4.5.3.56
openslide-python 1.1.2
opt-einsum 3.3.0
optuna 2.9.1
orderedmultidict 1.0.1
ortools 9.0.9048
osmnx 1.1.1
overrides 3.1.0
packaging 21.0
palettable 3.3.0
pandarallel 1.5.2
pandas 1.2.5
pandas-datareader 0.10.0
pandas-profiling 3.0.0
pandas-summary 0.0.7
pandasql 0.7.3
pandocfilters 1.4.2
panel 0.12.1
papermill 2.3.3
param 1.11.1
parso 0.8.2
partd 1.2.0
path 16.2.0
path.py 12.5.0
pathos 0.2.8
pathspec 0.9.0
pathtools 0.1.2
patsy 0.5.1
patternfork-nosql 3.6
pbr 5.6.0
pdf2image 1.16.0
pdfminer.six 20201018
PDPbox 0.2.1
pexpect 4.8.0
phik 0.11.2
pickleshare 0.7.5
Pillow 8.2.0
pip 21.2.4
plac 1.1.3
plotly 5.3.1
plotly-express 0.4.1
plotnine 0.8.0
pluggy 1.0.0
pointpats 2.2.0
polyglot 16.7.4
pooch 1.5.1
portalocker 2.3.2
portend 2.7.1
pox 0.3.0
poyo 0.5.0
ppca 0.0.4
ppft 1.6.6.4
preprocessing 0.1.13
preshed 3.0.5
prettytable 2.1.0
prometheus-client 0.11.0
promise 2.3
prompt-toolkit 3.0.19
pronouncing 0.2.0
proto-plus 1.19.0
protobuf 3.18.0
psutil 5.8.0
ptyprocess 0.7.0
pudb 2021.1
PuLP 2.5.0
py 1.10.0
py-lz4framed 0.14.0
py-stringmatching 0.4.2
py-stringsimjoin 0.3.2
pyaml 21.8.3
PyArabic 0.6.14
pyarrow 4.0.1
pyasn1 0.4.8
pyasn1-modules 0.2.7
PyAstronomy 0.16.0
pybind11 2.7.1
pycodestyle 2.7.0
pycosat 0.6.3
pycountry 20.7.3
pycparser 2.20
pycrypto 2.6.1
pyct 0.4.8
pycuda 2021.1
pydantic 1.8.2
pydash 5.0.2
pydegensac 0.1.2
pyDeprecate 0.3.1
pydicom 2.2.1
pydot 1.4.2
pydub 0.25.1
pyemd 0.5.1
pyexcel-io 0.6.4
pyexcel-ods 0.6.0
pyfasttext 0.4.6
pyflakes 2.3.1
pygeos 0.10.2
Pygments 2.10.0
PyJWT 2.1.0
pykalman 0.9.5
pyLDAvis 3.2.2
pymc3 3.11.4
PyMeeus 0.5.11
pymongo 3.12.0
Pympler 0.9
pynndescent 0.5.4
pynvml 8.0.4
pynvrtc 9.2
pyocr 0.8
pyOpenSSL 20.0.1
pyparsing 2.4.7
pyPdf 1.13
pyperclip 1.8.2
PyPrind 2.11.3
pyproj 3.1.0
pyrsistent 0.17.3
pysal 2.5.0
pyshp 2.1.3
PySocks 1.7.1
pystan 2.19.1.1
pytesseract 0.3.8
pytest 6.2.5
pytest-cov 2.12.1
python-bidi 0.4.2
python-dateutil 2.8.0
python-docx 0.8.11
python-igraph 0.9.6
python-Levenshtein 0.12.2
python-louvain 0.15
python-slugify 5.0.2
pytools 2021.2.8
pytorch-ignite 0.4.6
pytorch-lightning 1.4.4
pytz 2021.1
pyu2f 0.1.5
PyUpSet 0.1.1.post7
pyviz-comms 2.1.0
PyWavelets 1.1.1
PyYAML 5.4.1
pyzmq 22.2.1
qgrid 1.3.1
qtconsole 5.1.1
QtPy 1.11.1
quantecon 0.5.1
quilt3 3.5.0
randomgen 1.16.6
rasterio 1.2.8
rasterstats 0.15.0
ray 1.6.0
redis 3.5.3
regex 2021.8.3
requests 2.25.1
requests-futures 1.0.0
requests-oauthlib 1.3.0
requests-unixsocket 0.2.0
resampy 0.2.2
retrying 1.3.3
rgf-python 3.11.0
rmm 21.8.2
rsa 4.7.2
Rtree 0.9.7
ruamel-yaml-conda 0.15.100
rvlib 0.0.5
s2sphere 0.2.5
s3fs 2021.8.1
s3transfer 0.5.0
sacremoses 0.0.45
scattertext 0.1.4
scikit-image 0.18.3
scikit-learn 0.23.2
scikit-multilearn 0.2.0
scikit-optimize 0.8.1
scikit-plot 0.3.7
scikit-surprise 1.1.1
scipy 1.7.1
seaborn 0.11.2
segregation 2.1.0
semver 2.13.0
Send2Trash 1.8.0
sentencepiece 0.1.96
sentry-sdk 1.3.1
setuptools 57.4.0
setuptools-git 1.2
sgmllib3k 1.0.0
shap 0.39.0
Shapely 1.7.1
shortuuid 1.0.1
SimpleITK 2.1.1
simplejson 3.17.3
six 1.15.0
sklearn 0.0
sklearn-contrib-py-earth 0.1.0+1.gdde5f89
sklearn-pandas 2.2.0
slicer 0.0.7
smart-open 5.2.1
smhasher 0.150.1
smmap 3.0.5
sniffio 1.2.0
snuggs 1.4.7
sortedcontainers 2.4.0
SoundFile 0.10.3.post1
soupsieve 2.2.1
spacy 2.3.7
spaghetti 1.6.2
spectral 0.22.2
spglm 1.0.8
sphinx-rtd-theme 0.2.4
spint 1.0.7
splot 1.1.4
spopt 0.1.2
spreg 1.2.4
spvcm 0.3.0
SQLAlchemy 1.4.22
sqlitedict 1.7.0
sqlparse 0.4.1
squarify 0.4.3
srsly 1.0.5
statsmodels 0.12.2
stemming 1.0.1
stevedore 3.4.0
stop-words 2018.7.23
stopit 1.1.2
stumpy 1.9.2
subprocess32 3.5.4
sympy 1.8
tables 3.6.1
tabulate 0.8.9
tangled-up-in-unicode 0.1.0
tblib 1.7.0
tempora 4.1.1
tenacity 8.0.1
tensorboard 2.6.0
tensorboard-data-server 0.6.1
tensorboard-plugin-wit 1.8.0
tensorboardX 2.2
tensorflow 2.4.1
tensorflow-addons 0.12.1
tensorflow-cloud 0.1.13
tensorflow-datasets 3.0.0
tensorflow-estimator 2.4.0
tensorflow-gcs-config 2.4.0
tensorflow-hub 0.12.0
tensorflow-metadata 1.2.0
tensorflow-probability 0.14.1
tensorpack 0.11
termcolor 1.1.0
terminado 0.10.1
testpath 0.5.0
text-unidecode 1.3
textblob 0.15.3
texttable 1.6.4
textwrap3 0.9.2
Theano 1.0.5
Theano-PyMC 1.1.2
thinc 7.4.5
threadpoolctl 2.2.0
tifffile 2021.8.30
tobler 0.8.2
tokenizers 0.10.3
toml 0.10.2
tomli 1.2.1
toolz 0.11.1
torch 1.7.1+cu110
torchaudio 0.7.2
torchmetrics 0.5.0
torchtext 0.8.1
torchvision 0.8.2+cu110
tornado 6.1
TPOT 0.11.7
tqdm 4.62.1
traitlets 5.0.5
traittypes 0.2.1
transformers 4.5.1
treelite 2.0.0
treelite-runtime 2.0.0
trueskill 0.4.5
tsfresh 0.18.0
Twisted 21.7.0
txaio 21.2.1
typed-ast 1.4.3
typeguard 2.12.1
typing-extensions 3.7.4.3
typing-inspect 0.7.1
tzlocal 3.0
ucx-py 0.21.0
ujson 4.0.2
umap-learn 0.5.1
Unidecode 1.2.0
update-checker 0.18.0
uritemplate 3.0.1
urllib3 1.26.6
urwid 2.1.2
vecstack 0.4.0
visions 0.7.1
vowpalwabbit 8.11.0
vtk 9.0.3
Wand 0.6.7
wandb 0.12.2
wasabi 0.8.2
wavio 0.0.4
wcwidth 0.2.5
webencodings 0.5.1
websocket-client 0.57.0
Werkzeug 2.0.1
wfdb 3.4.1
wheel 0.37.0
whichcraft 0.6.1
widgetsnbextension 3.5.1
Wordbatch 1.4.8
wordcloud 1.8.1
wordsegment 1.3.1
wrapt 1.12.1
wslink 1.0.7
xarray 0.19.0
xgboost 1.4.2
xvfbwrapper 0.2.9
xxhash 2.0.2
yacs 0.1.8
yarl 1.6.3
yaspin 2.1.0
yellowbrick 1.3.post1
zc.lockfile 2.0
zict 2.0.0
zipp 3.5.0
zope.interface 5.4.0

The code:

time_series_with_nans = [-1., 1., np.nan, 2.4, np.nan, 5]
observed_time_series = tfp.sts.MaskedTimeSeries(
  time_series=time_series_with_nans,
  is_missing=tf.math.is_nan(time_series_with_nans))

# Build model using observed time series to set heuristic priors.
linear_trend_model = tfp.sts.LocalLinearTrend(
  observed_time_series=observed_time_series)
model = tfp.sts.Sum([linear_trend_model],
                    observed_time_series=observed_time_series)

# Fit model to data
parameter_samples, _ = tfp.sts.fit_with_hmc(model, observed_time_series)

The error:

Start by reproducing the TensorFlow Probability sts.impute_missing_values example linked in the report with the listed package versions. Trace the build_factored_surrogate_posterior() call and identify why batch_shape is rejected; done means the example runs without that unexpected-keyword error.