Keras graph visualization fails for models using mixed_float16 policy

[bersbersbers]

Diagnostics

Diagnostics output

--- check: autoidentify
INFO: diagnose_tensorboard.py version e43767ef2b648d0d5d57c00f38ccbd38390e38da

--- check: general
INFO: sys.version_info: sys.version_info(major=3, minor=9, micro=10, releaselevel='final', serial=0)
INFO: os.name: posix
INFO: os.uname(): posix.uname_result(sysname='Linux', nodename='xxx', release='xxx', version='xxx', machine='x86_64')
INFO: sys.getwindowsversion(): N/A

--- check: package_management
INFO: has conda-meta: False
INFO: $VIRTUAL_ENV: None

--- check: installed_packages
INFO: installed: tensorboard==2.8.0
INFO: installed: tensorflow==2.8.0rc1
INFO: installed: tf-estimator-nightly==2.8.0.dev2021122109
INFO: installed: tensorboard-data-server==0.6.1

--- check: tensorboard_python_version
INFO: tensorboard.version.VERSION: '2.8.0'

--- check: tensorflow_python_version
INFO: tensorflow.__version__: '2.8.0-rc1'
INFO: tensorflow.__git_version__: 'v2.8.0-rc0-49-g244b9d77fd4'

--- check: tensorboard_data_server_version
INFO: data server binary: '/opt/pyenv/versions/3.9.10/lib/python3.9/site-packages/tensorboard_data_server/bin/server'
INFO: data server binary version: b'rustboard 0.6.1'

--- check: tensorboard_binary_path
INFO: which tensorboard: b'/opt/pyenv/versions/3.9.10/bin/tensorboard\n'

--- check: addrinfos
socket.has_ipv6 = True
socket.AF_UNSPEC = <AddressFamily.AF_UNSPEC: 0>
socket.SOCK_STREAM = <SocketKind.SOCK_STREAM: 1>
socket.AI_ADDRCONFIG = <AddressInfo.AI_ADDRCONFIG: 32>
socket.AI_PASSIVE = <AddressInfo.AI_PASSIVE: 1>
Loopback flags: <AddressInfo.AI_ADDRCONFIG: 32>
Loopback infos: [(<AddressFamily.AF_INET6: 10>, <SocketKind.SOCK_STREAM: 1>, 6, '', ('::1', 0, 0, 0)), (<AddressFamily.AF_INET: 2>, <SocketKind.SOCK_STREAM: 1>, 6, '', ('127.0.0.1', 0))]
Wildcard flags: <AddressInfo.AI_PASSIVE: 1>
Wildcard infos: [(<AddressFamily.AF_INET: 2>, <SocketKind.SOCK_STREAM: 1>, 6, '', ('0.0.0.0', 0)), (<AddressFamily.AF_INET6: 10>, <SocketKind.SOCK_STREAM: 1>, 6, '', ('::', 0, 0, 0))]

--- check: readable_fqdn
INFO: socket.getfqdn(): 'xxx'

--- check: stat_tensorboardinfo
INFO: directory: /tmp/.tensorboard-info
INFO: os.stat(...): os.stat_result(st_mode=16895, st_ino=3449412, st_dev=66306, st_nlink=2, st_uid=145278, st_gid=890, st_size=4096, st_atime=1643746598, st_mtime=1643759711, st_ctime=1643759711)
INFO: mode: 0o40777

--- check: source_trees_without_genfiles
INFO: tensorboard_roots (1): ['/opt/pyenv/versions/3.9.10/lib/python3.9/site-packages']; bad_roots (0): []

--- check: full_pip_freeze
INFO: pip freeze --all:
absl-py==1.0.0
alabaster==0.7.12
appdirs==1.4.4
astroid==2.9.3
astunparse==1.6.3
attrs==21.4.0
Babel==2.9.1
bandit==1.7.2
black==22.1.0
blessed==1.19.1
cachetools==5.0.0
certifi==2021.10.8
cffconvert==2.0.0
charset-normalizer==2.0.11
click==8.0.3
colour-science==0.3.16
cycler==0.11.0
Deprecated==1.2.13
docopt==0.6.2
docutils==0.17.1
eradicate==2.0.0
ExifRead==2.3.2
flake8==4.0.1
flake8-black==0.2.4
flake8-eradicate==1.2.0
flake8-plugin-utils==1.3.2
flake8-pytest-style==1.6.0
flatbuffers==2.0
flyingcircus==0.1.4.0
flyingcircus-numeric==0.1.1.2
fonttools==4.29.0
gast==0.5.3
gitdb==4.0.9
GitPython==3.1.26
google-auth==2.5.0
google-auth-oauthlib==0.4.6
google-pasta==0.2.0
grpcio==1.43.0
h5py==3.6.0
humanize==3.14.0
idna==3.3
imageio==2.14.1
imagesize==1.3.0
importlib-metadata==4.10.1
iniconfig==1.1.1
isort==5.10.1
Jinja2==3.0.3
joblib==1.1.0
json5==0.9.6
jsonschema==3.2.0
keras==2.8.0rc1
Keras-Preprocessing==1.1.2
kiwisolver==1.3.2
lazy-object-proxy==1.7.1
libclang==13.0.0
lxml==4.7.1
Markdown==3.3.6
markdown2==2.4.2
MarkupSafe==2.0.1
matplotlib==3.5.1
mccabe==0.6.1
mypy==0.931
mypy-extensions==0.4.3
natsort==8.1.0
networkx==2.6.3
numpy==1.22.1
oauthlib==3.2.0
opencv-python==4.5.5.62
opt-einsum==3.3.0
ordered-set==4.1.0
packaging==21.3
pandas==1.4.0
pathspec==0.9.0
pbr==5.8.0
Pillow==9.0.0
pip==22.0.2
pipdeptree==2.2.1
platformdirs==2.4.1
pluggy==1.0.0
protobuf==3.19.4
psutil==5.9.0
py==1.11.0
pyasn1==0.4.8
pyasn1-modules==0.2.8
pycodestyle==2.8.0
pydocstyle==6.1.1
pyflakes==2.4.0
Pygments==2.11.2
pykwalify==1.8.0
pylama==8.3.7
pylint==2.12.2
pynrrd==0.4.2
pyparsing==3.0.7
pyrsistent==0.18.1
pytest==6.2.5
python-dateutil==2.8.2
python-pptx==0.6.21
pytz==2021.3
PyWavelets==1.2.0
PyYAML==6.0
raster-geometry==0.1.4.1
requests==2.27.1
requests-oauthlib==1.3.1
rsa==4.8
ruamel.yaml==0.17.20
ruamel.yaml.clib==0.2.6
scikit-image==0.19.1
scikit-learn==1.0.2
scikit-plot==0.3.7
scipy==1.7.3
setuptools==60.6.0
setuptools-scm==6.4.2
SharedArray==3.2.1
silence-tensorflow==1.1.1
six==1.16.0
smmap==5.0.0
snowballstemmer==2.2.0
Sphinx==4.4.0
sphinxcontrib-applehelp==1.0.2
sphinxcontrib-devhelp==1.0.2
sphinxcontrib-htmlhelp==2.0.0
sphinxcontrib-jsmath==1.0.1
sphinxcontrib-qthelp==1.0.3
sphinxcontrib-serializinghtml==1.1.5
stevedore==3.5.0
tensorboard==2.8.0
tensorboard-data-server==0.6.1
tensorboard-plugin-wit==1.8.1
tensorflow==2.8.0rc1
tensorflow-addons==0.15.0
tensorflow-io-gcs-filesystem==0.23.1
termcolor==1.1.0
tf-estimator-nightly==2.8.0.dev2021122109
tf-explain==0.3.1
tf-keras-vis==0.8.0
threadpoolctl==3.0.0
tifffile==2021.11.2
toml==0.10.2
tomli==2.0.0
toolz==0.11.2
typeguard==2.13.3
types-requests==2.27.7
types-urllib3==1.26.8
typing_extensions==4.0.1
urllib3==1.26.8
wcwidth==0.2.5
Werkzeug==2.0.2
wheel==0.37.1
wrapt==1.13.3
XlsxWriter==3.0.2
zipp==3.7.0

Issue description

• Run code below
• Run tensorboard --logdir /tmp/bug
• Open http://localhost:6006/#graphs&run=train
• Choose Tag (2) == keras
• Boom.

"""Demonstrate bug."""
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

from tensorflow import keras
keras.mixed_precision.set_global_policy("mixed_float16")  # works fine without this line

model = keras.Sequential([keras.layers.Dense(1, input_shape=(1,))])
callback = keras.callbacks.TensorBoard("/tmp/bug")
model.compile(loss="mse", optimizer="sgd")
model.fit(x=[0], y=[0], callbacks=[callback])

Log:

E0202 00:50:31.522159 139930501322496 _internal.py:225] Error on request:
Traceback (most recent call last):
  File "/opt/pyenv/versions/3.9.10/lib/python3.9/site-packages/werkzeug/serving.py", line 319, in run_wsgi
    execute(self.server.app)
  File "/opt/pyenv/versions/3.9.10/lib/python3.9/site-packages/werkzeug/serving.py", line 308, in execute
    application_iter = app(environ, start_response)
  File "/opt/pyenv/versions/3.9.10/lib/python3.9/site-packages/tensorboard/backend/application.py", line 525, in __call__
    return self._app(environ, start_response)
  File "/opt/pyenv/versions/3.9.10/lib/python3.9/site-packages/tensorboard/backend/application.py", line 566, in wrapper
    return wsgi_app(environ, start_response)
  File "/opt/pyenv/versions/3.9.10/lib/python3.9/site-packages/tensorboard/backend/security_validator.py", line 77, in __call__
    return self._application(environ, start_response_proxy)
  File "/opt/pyenv/versions/3.9.10/lib/python3.9/site-packages/tensorboard/backend/path_prefix.py", line 68, in __call__
    return self._application(environ, start_response)
  File "/opt/pyenv/versions/3.9.10/lib/python3.9/site-packages/tensorboard/backend/experiment_id.py", line 73, in __call__
    return self._application(environ, start_response)
  File "/opt/pyenv/versions/3.9.10/lib/python3.9/site-packages/tensorboard/backend/empty_path_redirect.py", line 43, in __call__
    return self._application(environ, start_response)
  File "/opt/pyenv/versions/3.9.10/lib/python3.9/site-packages/tensorboard/backend/application.py", line 589, in wrapper
    return wsgi_app(environ, start_response)
  File "/opt/pyenv/versions/3.9.10/lib/python3.9/site-packages/tensorboard/backend/application.py", line 548, in _route_request
    return self.exact_routes[clean_path](environ, start_response)
  File "/opt/pyenv/versions/3.9.10/lib/python3.9/site-packages/werkzeug/wrappers/request.py", line 206, in application
    resp = f(*args[:-2] + (request,))
  File "/opt/pyenv/versions/3.9.10/lib/python3.9/site-packages/tensorboard/plugins/graph/graphs_plugin.py", line 307, in graph_route
    result = self.graph_impl(
  File "/opt/pyenv/versions/3.9.10/lib/python3.9/site-packages/tensorboard/plugins/graph/graphs_plugin.py", line 222, in graph_impl
    graph = keras_util.keras_model_to_graph_def(keras_model_config)
  File "/opt/pyenv/versions/3.9.10/lib/python3.9/site-packages/tensorboard/plugins/graph/keras_util.py", line 238, in keras_model_to_graph_def
    tf_dtype = dtypes.as_dtype(layer_config.get("dtype"))
  File "/opt/pyenv/versions/3.9.10/lib/python3.9/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py", line 659, in as_dtype
    return _INTERN_TABLE[type_value]
TypeError: unhashable type: 'dict'

Debugging:

type_value
{'class_name': 'Policy', 'config': {'name': 'mixed_float16'}, 'shared_object_id': 3}

[japie1235813]

Thanks for filing the bug. Have verified that the bug is reproducible.

[pindinagesh]

@bersbersbers

Can you please confirm if we can close this issue since PR is merged. Thanks!

[bersbersbers]

@pindinagesh I haven't tested it, but I guess the PR will fix it. Thanks!