如何在 Keras TimeDistributed 层中使用 tfp.layers.IndependentNormal？

Question

该模型运行良好，没有任何错误。

import tensorflow as tf
from keras import Sequential
from keras.layers import Bidirectional, LSTM, Dropout, Dense, TimeDistributed
import tensorflow_probability as tfp


model = Sequential()
model.add(Bidirectional(LSTM(16, return_sequences=True), input_shape=(100, 1), ))
model.add(Dropout(0.1))
model.add(LSTM(16, return_sequences=True))
model.add(Dropout(0.1))
model.add(LSTM(16, return_sequences=True))
model.add(Dropout(0.1))
model.add(TimeDistributed(Dense(1)))

print(model.summary())

输出：

Model: "sequential"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
bidirectional (Bidirectional (None, 100, 32)           2304      
_________________________________________________________________
dropout (Dropout)            (None, 100, 32)           0         
_________________________________________________________________
lstm_1 (LSTM)                (None, 100, 16)           3136      
_________________________________________________________________
dropout_1 (Dropout)          (None, 100, 16)           0         
_________________________________________________________________
lstm_2 (LSTM)                (None, 100, 16)           2112      
_________________________________________________________________
dropout_2 (Dropout)          (None, 100, 16)           0         
_________________________________________________________________
time_distributed (TimeDistri (None, 100, 1)            17        
=================================================================
Total params: 7,569
Trainable params: 7,569
Non-trainable params: 0
_________________________________________________________________
None

但是在我将 model.add(TimeDistributed(Dense(1))) 行更改为 model.add(TimeDistributed(tfp.layers.IndependentNormal(1)))，该模型运行时出错。

错误是：

Traceback (most recent call last):
  File "/Users/xxx.py", line 27, in <module>
    model.add(TimeDistributed(tfp.layers.IndependentNormal(1)))
  File "/Library/Frameworks/Python.framework/Versions/3.6/lib/python3.6/site-packages/tensorflow/python/training/tracking/base.py", line 530, in _method_wrapper
    result = method(self, *args, **kwargs)
  File "/Library/Frameworks/Python.framework/Versions/3.6/lib/python3.6/site-packages/keras/engine/sequential.py", line 217, in add
    output_tensor = layer(self.outputs[0])
  File "/Library/Frameworks/Python.framework/Versions/3.6/lib/python3.6/site-packages/keras/engine/base_layer.py", line 977, in __call__
    input_list)
  File "/Library/Frameworks/Python.framework/Versions/3.6/lib/python3.6/site-packages/keras/engine/base_layer.py", line 1115, in _functional_construction_call
    inputs, input_masks, args, kwargs)
  File "/Library/Frameworks/Python.framework/Versions/3.6/lib/python3.6/site-packages/keras/engine/base_layer.py", line 848, in _keras_tensor_symbolic_call
    return self._infer_output_signature(inputs, args, kwargs, input_masks)
  File "/Library/Frameworks/Python.framework/Versions/3.6/lib/python3.6/site-packages/keras/engine/base_layer.py", line 888, in _infer_output_signature
    outputs = call_fn(inputs, *args, **kwargs)
  File "/Library/Frameworks/Python.framework/Versions/3.6/lib/python3.6/site-packages/keras/layers/wrappers.py", line 271, in call
    output_shape = self.compute_output_shape(input_shape)
  File "/Library/Frameworks/Python.framework/Versions/3.6/lib/python3.6/site-packages/keras/layers/wrappers.py", line 190, in compute_output_shape
    child_output_shape = self.layer.compute_output_shape(child_input_shape)
  File "/Library/Frameworks/Python.framework/Versions/3.6/lib/python3.6/site-packages/keras/utils/tf_utils.py", line 259, in wrapper
    output_shape = fn(instance, input_shape)
  File "/Library/Frameworks/Python.framework/Versions/3.6/lib/python3.6/site-packages/keras/layers/core.py", line 867, in compute_output_shape
    return super(Lambda, self).compute_output_shape(input_shape)
  File "/Library/Frameworks/Python.framework/Versions/3.6/lib/python3.6/site-packages/keras/engine/base_layer.py", line 790, in compute_output_shape
    return tf.nest.map_structure(lambda t: t.shape, outputs)
  File "/Library/Frameworks/Python.framework/Versions/3.6/lib/python3.6/site-packages/tensorflow/python/util/nest.py", line 869, in map_structure
    structure[0], [func(*x) for x in entries],
  File "/Library/Frameworks/Python.framework/Versions/3.6/lib/python3.6/site-packages/tensorflow/python/util/nest.py", line 869, in <listcomp>
    structure[0], [func(*x) for x in entries],
  File "/Library/Frameworks/Python.framework/Versions/3.6/lib/python3.6/site-packages/keras/engine/base_layer.py", line 790, in <lambda>
    return tf.nest.map_structure(lambda t: t.shape, outputs)
  File "/Library/Frameworks/Python.framework/Versions/3.6/lib/python3.6/site-packages/keras/engine/keras_tensor.py", line 131, in shape
    return self._type_spec._shape  # pylint: disable=protected-access
AttributeError: 'UserRegisteredSpec' object has no attribute '_shape'

那么如何在Keras TimeDistributed层中使用tfp.layers.IndependentNormal层呢？

非常感谢你对我的帮助。

=================================================== ==========

本项目使用的包版本：

Keras: 2.6.0

TensorFlow: 2.6.2

TensorFlow-probability: 0.14.1