AttributeError: 'SymbolicTensor' object has no attribute 'log_prob' when exporting train signature with `IndependentNormal` layer

[AndiBerber]

I want to export the training signature to train my model with the C++ API. However, I am not able to export the model, even after reading through [https://github.com/tensorflow/probability/issues/742]() , https://github.com/tensorflow/tensorflow/issues/36181 [https://stackoverflow.com/questions/59743872/when-training-a-variational-bayesian-neural-network-in-tfp-how-can-i-visualize]() My conda list | grep tensorflow versions are: tensorflow 2.15.0.post1
tensorflow-estimator 2.15.0
tensorflow-io-gcs-filesystem 0.36.0
tensorflow-probability 0.23.0

My conda list | grep keras versions are: keras 2.15.0
keras-preprocessing 1.1.2

import numpy as np
from tensorflow.keras.layers import Input, Dense
import tensorflow_probability as tfp
from tensorflow.keras.models import Model
from tensorflow.keras.optimizers import Adam
import tensorflow as tf

inputDim = 10
targetDim = 1

#build train data
samples = 1000

input_list = []
for ii in range(samples):
    input_list.append(np.arange(inputDim))

input_arr = np.array(input_list)
target = arr = np.random.normal(5.0, 0.5, (samples,1))

#define model
input = Input(shape=(inputDim))
distribution_params = Dense(2)(input)
outputs = tfp.layers.IndependentNormal(targetDim)(distribution_params)

#define loss
def nll(targets, estimated_distribution):
    return -estimated_distribution.log_prob(targets)

#compile and fit model
optimizer = Adam()
model = Model(inputs= [input] , outputs=[outputs])
model.compile(optimizer=optimizer, loss=nll)#, metrics = lossFunction)
model.summary()
model.fit(input_arr,target, shuffle=True, epochs=500)#, verbose = 2)

# test prediction
prediction = model(np.expand_dims(np.arange(inputDim), axis = 0))
print("prediction mean : ", prediction.mean())
print("stdDev = ", prediction.stddev())

#export training signature
@tf.function
def trainOp(inputs, targets):
    ### has to return loss ###
    with tf.GradientTape() as tape:
        predictions  = model(inputs)
        loss = nll(predictions, targets)
    gradients = tape.gradient(loss, model.trainable_variables)    
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    return loss

signatures = {}
signatures["trainOp"] = trainOp.get_concrete_function(inputs = tf.TensorSpec((None, inputDim), tf.float32), 
                                                      targets = tf.TensorSpec((None, targetDim), tf.float32))

model.save('./testExport/', save_traces = False, signatures = signatures)

failes with : AttributeError: in user code:

File "/home/aberberich/Shared/Andi/tf2Api/reworked/min_export_failure.py", line 57, in trainOp  *
    loss = nll(predictions, targets)
File "/home/aberberich/Shared/Andi/tf2Api/reworked/min_export_failure.py", line 37, in nll  *
    return -estimated_distribution.log_prob(targets)

AttributeError: 'SymbolicTensor' object has no attribute 'log_prob'`

[AndiBerber]

loss = nll(predictions, targets)

should be

loss = nll(targets, predictions)