Propagation with multiple weight samples in denseVariational not working

Currently the forward propagation when using multiple weight samples at once is not possible in DenseVariational. There may be many situations where this functionality is desirable (for instance, it is necessary when using alpha-divergences as loss function instead of variational Bayes).

Tensorflow: 2.1.0
tfp: 0.90.

We can specify the convert_to_tensor function in DistributionLambda to sample multiple times, e.g:

def posterior_mean_field(self, kernel_size, bias_size=0, n_samples=25, dtype=None):
    n_parameter = kernel_size + bias_size
    c_init = np.log(np.expm1(0.2))
    return tf.keras.Sequential(
        [
            tfp.layers.VariableLayer(2 * n_parameter, dtype=dtype),
            tfp.layers.DistributionLambda(
                lambda t: tfd.Independent(
                    tfd.Normal(
                        loc=t[..., :n_parameter],
                        scale=1e-5 + tf.nn.softplus(c_init + t[..., n_parameter:]),
                    ),
                    reinterpreted_batch_ndims=1,
                ),
                convert_to_tensor_fn=lambda dist: dist.sample(n_samples),
            ),
        ]
    )

Then everything will work as expected until the call method() of dense_variational fails here because

      outputs = tf.nn.bias_add(outputs, bias)

expects a 2-D Array and not 3-D tensors. The fix seems to be straightforward, not using tf.nn.bias_add here, but writing the bias addition explicitly with the possibility of having multiple dimensions.

My current workaround is defining a custom layer that does the bias addition explicitly:

class DenseVariationalFix(tf.keras.Model):
    def __init__(self, *args, activation=None, **kwargs):
        super().__init__()
        self.full_layer = tfp.layers.DenseVariational(
            *args, activation=None, use_bias=False,  **kwargs
        )
        self.bias_layer = tfp.layers.DenseVariational(
            *args, activation=None, use_bias=False,  **kwargs
        )

        self.activation = activation

    def call(self, inputs):
        lhs = self.full_layer(inputs)
        rhs = self.bias_layer(tf.ones_like(inputs[..., :1]))
        addit = tf.keras.layers.Lambda(lambda inputs: inputs[0] + inputs[1])(
            [lhs, rhs]
        )
        return tf.keras.layers.Activation(self.activation)(addit)

Start in tensorflow_probability/python/layers/dense_variational_v2.py at the call method and reproduce the issue with the provided multi-sample DistributionLambda setup. Trace the bias addition for the resulting higher-rank tensor; done means DenseVariational supports multiple weight samples during forward propagation without the shown shape failure.