Problem using Variational Dense from tensorflow probability

Hello,

I am trying to build a bayesian neural network, and am running into the following Error, when I use the VariationalDense layer. I should note that this works fine using regular Dense from keras.layers.

Thanks!

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-35-7c804e9ec8ee> in <module>
----> 1 dc.build_architecture(fclayer=fclayer,
      2                       n_layers=nlayers,
      3                       latent_dim=latent_dim,
      4                       l2_weight=l2_weight)
      5 

<ipython-input-32-a2070c82d9f0> in build_architecture(self, fclayer, n_layers, latent_dim, l2_weight)
     94         # FC Layer
     95         print(inputs.shape)
---> 96         x = tfp.layers.DenseVariational(fclayer, 
     97                                         make_prior_fn=prior,
     98                                         make_posterior_fn=posterior,

~/project-rlittman/miniconda3/envs/tf_pb/lib/python3.9/site-packages/tensorflow/python/keras/engine/base_layer.py in __call__(self, *args, **kwargs)
    949     # >> model = tf.keras.Model(inputs, outputs)
    950     if _in_functional_construction_mode(self, inputs, args, kwargs, input_list):
--> 951       return self._functional_construction_call(inputs, args, kwargs,
    952                                                 input_list)
    953 

~/project-rlittman/miniconda3/envs/tf_pb/lib/python3.9/site-packages/tensorflow/python/keras/engine/base_layer.py in _functional_construction_call(self, inputs, args, kwargs, input_list)
   1088           layer=self, inputs=inputs, build_graph=True, training=training_value):
   1089         # Check input assumptions set after layer building, e.g. input shape.
-> 1090         outputs = self._keras_tensor_symbolic_call(
   1091             inputs, input_masks, args, kwargs)
   1092 

~/project-rlittman/miniconda3/envs/tf_pb/lib/python3.9/site-packages/tensorflow/python/keras/engine/base_layer.py in _keras_tensor_symbolic_call(self, inputs, input_masks, args, kwargs)
    820       return nest.map_structure(keras_tensor.KerasTensor, output_signature)
    821     else:
--> 822       return self._infer_output_signature(inputs, args, kwargs, input_masks)
    823 
    824   def _infer_output_signature(self, inputs, args, kwargs, input_masks):

~/project-rlittman/miniconda3/envs/tf_pb/lib/python3.9/site-packages/tensorflow/python/keras/engine/base_layer.py in _infer_output_signature(self, inputs, args, kwargs, input_masks)
    861           # TODO(kaftan): do we maybe_build here, or have we already done it?
    862           self._maybe_build(inputs)
--> 863           outputs = call_fn(inputs, *args, **kwargs)
    864 
    865         self._handle_activity_regularization(inputs, outputs)

~/project-rlittman/miniconda3/envs/tf_pb/lib/python3.9/site-packages/tensorflow/python/autograph/impl/api.py in wrapper(*args, **kwargs)
    665       try:
    666         with conversion_ctx:
--> 667           return converted_call(f, args, kwargs, options=options)
    668       except Exception as e:  # pylint:disable=broad-except
    669         if hasattr(e, 'ag_error_metadata'):

~/project-rlittman/miniconda3/envs/tf_pb/lib/python3.9/site-packages/tensorflow/python/autograph/impl/api.py in converted_call(f, args, kwargs, caller_fn_scope, options)
    348   if conversion.is_in_allowlist_cache(f, options):
    349     logging.log(2, 'Allowlisted %s: from cache', f)
--> 350     return _call_unconverted(f, args, kwargs, options, False)
    351 
    352   if ag_ctx.control_status_ctx().status == ag_ctx.Status.DISABLED:

~/project-rlittman/miniconda3/envs/tf_pb/lib/python3.9/site-packages/tensorflow/python/autograph/impl/api.py in _call_unconverted(f, args, kwargs, options, update_cache)
    476 
    477   if kwargs is not None:
--> 478     return f(*args, **kwargs)
    479   return f(*args)
    480 

Start with the tfp.layers.DenseVariational entry point and reproduce the provided notebook call using the shown prior and posterior arguments. The issue does not name a repository file, test, TensorFlow Probability version, or the actual ValueError message, so the failure cannot be localized or given a clear done condition from the report alone.