Get mean and std from probability-layer in a model with coupled features

Hallo,

I have coupled features like following and in the end of the network a probability-layer, such that I can get the mean and std of the model:

coupled_feature = [model1(input_layers[0]), input_layers[1]]
coupled_feature = keras.layers.Multiply()(coupled_feature)

prediction = layers.Dense(2, activation='relu', name='SumLayer')(coupled_feature)
prediction = tfp.layers.DistributionLambda(lambda x: tfd.Normal(loc=x[..., :1], scale=1e-3 + tf.math.softplus(-0.3 * x[..., 1:])), name='distributionLayer')(prediction)   #(lambda x: tfd.Normal(loc=x, scale=0.15), name='distributionLayer')  #(lambda x: tfd.Normal(loc=x[..., :1], scale=1e-3 + tf.math.softplus(0.05 * x[..., 1:])))(prediction)    #loc=x, scale=0.01
model = keras.models.Model(inputs=input_layers, outputs=prediction)

optimizer = keras.optimizers.RMSprop(0.001) #Adam(learning_rate=0.0001)   #
negloglik = lambda y, p_y: -p_y.log_prob(y)

model.compile(loss=negloglik,#loss='mse',#
              optimizer=optimizer,
              metrics=['mae', 'mse'])

Now I can compile and train this model without problems, but how to calculate the mean and std from the model?

The following does not work:

yhat = model([test_dataset.pictocode, test_dataset.loc[:, ['mean_max', 'mean_medium', 'mean_low']]])
test_mean = yhat.mean()

It gives me the following error:

tensorflow.python.framework.errors_impl.InvalidArgumentError: In[0] is not a matrix. Instead it has shape [65] [Op:MatMul]

But predicting works in the same way without problems:

test_predictions = model.predict([test_dataset.pictocode, test_dataset.loc[:, ['mean_max', 'mean_medium', 'mean_low']]]).flatten()

Is this not possible with coupled features at the moment or do I something wrong?

Many thanks fore some adwises!

The issue names no repository file, test, or entry point. Start by reproducing the shown DistributionLambda model and its yhat.mean() error, then inspect the relevant TensorFlow Probability distribution-layer behavior. Done means establishing the supported way to obtain mean and standard deviation for this coupled-input model and documenting or correcting the behavior.