Compound Poisson Keras custom loss function

Question

I am trying to implement a custom loss function using Tensorflow as the negative loglikelihood of this expression (which is a compound Poisson-Gamma):

The first term (represented by the Dirac delta) refers to the case when z == 0, while the sum (which needs to be truncated at some point in the implementation as it goes to infinity) represents the product of the probability from a Gamma and a Poisson distribution. This is the tentative implementation in Tensorflow:

import tensorflow as tf
import tensorflow_probability as tfp
from functools import partial
tf.enable_eager_execution()
import numpy as np

def pois_gamma_compound_loss(y_true, y_pred):
    lambda_, alpha, beta = y_pred[:, 0], y_pred[:, 1], y_pred[:, 2]

    poisson_distr = tfp.distributions.Poisson(rate=lambda_)
    ijk_0 = (1.0, tf.zeros_like(y_true))
    c = lambda i, p: i < 4
    b = lambda i, p: (tf.add(i, 1), 
                      p + tf.math.multiply(x = poisson_distr.prob(tf.zeros_like(y_true) + i),
                                           y = tfp.distributions.Gamma(concentration=tf.math.multiply(x = alpha,
                                                                                                      y = tf.zeros_like(y_true) + i), 
                                                                       rate=beta).prob(y_tru)))

    ijk_final = tf.while_loop(c, b, ijk_0)
    batch_lik = tf.add(ijk_final[1], tf.math.exp(tf.multiply(lambda_, -1.0))) 
    return -tf.reduce_mean(batch_lik)


inputs = Input(shape=(39,))
x = Dense(4, activation='relu', kernel_initializer='random_uniform')(inputs)
x = Dense(4, activation='relu', kernel_initializer='random_uniform')(inputs)
x = Dense(6, activation='relu', kernel_initializer='random_uniform')(inputs)

lambda_ = Dense(1, activation="softmax", name="lambda", kernel_initializer='random_uniform')(x)
alpha = Dense(1, activation="softmax", name="alpha", kernel_initializer='random_uniform')(x)
beta = Dense(1, activation="softmax", name="beta", kernel_initializer='random_uniform')(x)
output_params = Concatenate(name="pvec", axis=1)([lambda_, alpha, beta])

model = Model(inputs, output_params)
model.compile(loss=pois_gamma_compound_loss, optimizer='adam')
model.fit(X_train, y_train, epochs=60, batch_size=20)