I am attempting to use the Gamma function from tfp in a custom Keras loss function using the `log_prob`

method, but the function always returns `nan`

when training starts.

I have tested the loss function and seems to work fine:

```
import tensorflow as tf
import tensorflow_probability as tfp
tf.enable_eager_execution()
def gamma_loss(y_true, alpha, beta):
gamma_distr = tfp.distributions.Gamma(concentration=alpha, rate=beta)
log_lik_gamma = gamma_distr.log_prob(y_true)
return -tf.reduce_mean(log_lik_gamma)
gamma_loss(100, 2, 2).numpy()
# 194.00854
```

The problem may be related to the parameters (`alpha`

and `beta`

) that I am passing to the function and that are produced by the final (custom) layer of the model I am using.
This is the full snippet:

```
import tensorflow as tf
from tensorflow.keras import backend as K
from tensorflow.keras.layers import Input, Dense, Layer, Concatenate
from tensorflow.keras.models import Model
from tensorflow.keras.initializers import glorot_normal
import tensorflow_probability as tfp
from sklearn.datasets import make_regression
class GammaLayer(Layer):
def __init__(self, output_dim, **kwargs):
self.output_dim = output_dim
super(GammaLayer, self).__init__(**kwargs)
def build(self, input_shape):
n_weight_rows = 4
self.kernel_2 = self.add_weight(name='kernel_2',
shape=(n_weight_rows, self.output_dim),
initializer=glorot_normal(),
trainable=True)
self.kernel_3 = self.add_weight(name='kernel_3',
shape=(n_weight_rows, self.output_dim),
initializer=glorot_normal(),
trainable=True)
self.bias_2 = self.add_weight(name='bias_2',
shape=(self.output_dim,),
initializer=glorot_normal(),
trainable=True)
self.bias_3 = self.add_weight(name='bias_3',
shape=(self.output_dim,),
initializer=glorot_normal(),
trainable=True)
super(GammaLayer, self).build(input_shape)
def call(self, x):
# Here i use softplus to make the parameters strictly positive
alpha = tf.math.softplus(K.dot(x, self.kernel_2) + self.bias_2)
beta = tf.math.softplus(K.dot(x, self.kernel_3) + self.bias_3)
return [alpha, beta]
def compute_output_shape(self, input_shape):
"""
The assumption is that the output is always one-dimensional
"""
return [(input_shape[0], self.output_dim), (input_shape[0], self.output_dim)]
def gamma_loss(y_true, y_pred):
alpha, beta = y_pred[0], y_pred[1]
gamma_distr = tfp.distributions.Gamma(concentration=alpha, rate=beta)
return -tf.reduce_mean(gamma_distr.log_prob(y_true))
X, y = make_regression(n_samples=1000, n_features=3, noise=0.1)
inputs = Input(shape=(3,))
x = Dense(6, activation='relu')(inputs)
x = Dense(4, activation='relu')(x)
x = GammaLayer(1, name='main_output')(x)
output_params = Concatenate(1, name="pvec")(x)
model = Model(inputs, output_params)
model.compile(loss=gamma_loss, optimizer='adam')
model.fit(X, y, epochs=30, batch_size=10) ```
```

Can you try adding an additional 1e-6 or so outside the softplus? For very negative values, softplus becomes quite close to zero.