Getting mean and covariance matrix for multivariate normal from keras model
Given that the covariance matrix has to be positive definite, the cholesky decomposition is a good way to solve this problem.
So the output of the network will be the mean vector mu and the upper triangular part of the cholesky matrix (denoted T here). The diagonal of this matrix must be positive elements (the diagonal of the covariance matrix are standard deviations):
p = y_train.shape[1] # dimension of the covariance matrix 
inputs = Input(shape=(6,))
layer1 = Dense(24, activation='relu')(inputs)
layer2 = Dense(12, activation='relu')(layer1)
mu = Dense(p, activation = &quot;linear&quot;)(layer1)
T1 = Dense(p, activation=&quot;exponential&quot;)(layer1)# diagonal of T
T2 = Dense((p*(p-1)/2), activation=&quot;linear&quot;)(layer1)
outputs = Concatenate()([mu, T1, T2]) 

Now let's define the loss function. Firstly, let's define the function that will extract the outputs of the network:
def mu_sigma(output):
    mu = output[0][0:p]
    T1 = output[0][p:2*p]
    T2 = output[0][2*p:]
    ones = tf.ones((p,p), dtype=tf.float32) 
    mask_a = tf.linalg.band_part(ones, 0, -1)  
    mask_b = tf.linalg.band_part(ones, 0, 0)  
    mask = tf.subtract(mask_a, mask_b) 
    zero = tf.constant(0, dtype=tf.float32)
    non_zero = tf.not_equal(mask, zero)
    indices = tf.where(non_zero)
    T2 = tf.sparse.SparseTensor(indices,T2,dense_shape=tf.cast((p,p),
         dtype=tf.int64))
    T2 = tf.sparse.to_dense(T2)
    T1 = tf.linalg.diag(T1)
    sigma = T1 + T2
    return mu, sigma

Now for the loss function:
from tensorflow_probability import distributions as tfd
def gnll_loss(y, pred):
    mu, sigma = mu_sigma(pred)
    gm = tfd.MultivariateNormalTriL(loc=mu, scale_tril=sigma)
    log_likelihood = gm.log_prob(y)          
    return - tf.math.reduce_sum(log_likelihood)