Beta Distribution WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.

I'm building a functionally defined TF neural network that predicts 2 outputs that are modeled by with 2 different TFP distributions. If both outputs are modeled as Normal, then it looks like this:

import numpy as np
import pandas as pd
import tensorflow as tf
from tensorflow.keras import layers, Model
import tensorflow_probability as tfp
from tensorflow_probability import distributions as tfd

zero_buffer = 1e-5
input_layer = layers.Input(2)
dense_layer = layers.Dense(64, activation='relu')(input_layer)
pre_dist_layer = layers.Dense(4)(dense_layer)
dist1 = tfp.layers.DistributionLambda(
    lambda t: tfd.Normal(loc=t[..., 0],
                         scale=zero_buffer + tf.nn.softplus(t[..., 1])))(pre_dist_layer)
dist2 = tfp.layers.DistributionLambda(
    lambda t: tfd.Normal(loc=t[..., 2],
                         scale=zero_buffer + tf.nn.softplus(t[..., 3])))(pre_dist_layer)
# Had to add layer to define shape because user-defined DistributionLambda doesn't have .shape for some reason
dist_shape1 = layers.Reshape((1,), input_shape=(1,))(dist1)
dist_shape2 = layers.Reshape((1,), input_shape=(1,))(dist2)

concat_layer = layers.Concatenate(axis=1)([dist_shape1, dist_shape2])
model = Model(input_layer, concat_layer)
optimizer = tf.keras.optimizers.Adam()
model.compile(loss='mse', optimizer=optimizer, metrics=['mae', 'mse'])

This builds and compiles the model fine. However, when I swap out the 2nd distribution layer dist2 with a tfd.Beta distribution like this

dist2 = tfp.layers.DistributionLambda(
    lambda t: tfd.Beta(concentration0=zero_buffer + tf.nn.softplus(t[..., 2]),
                       concentration1=zero_buffer + tf.nn.softplus(t[..., 3])))(pre_dist_layer)

then I get the following warning:

WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.
WARNING:tensorflow:@custom_gradient grad_fn has 'variables' in signature, but no ResourceVariables were used on the forward pass.

I've tried doing the same with a couple other TFP distributions (e.g. tfd.LogNormal, tfd.GeneralizedExtremeValue) and I don't get this warning.

If it was just a warning, and the model still trained fine, I guess I'd be OK with that. However, when I use tfd.Beta and try to train the model eventually the loss and all the metrics turn to nan. Is there a bug in the tfd.Beta class causing this?

I'm using

• python==3.8.11
• tensorflow==2.10.0
• tensorflow-probability==0.18.0

Reproduce the supplied functional Keras model with DistributionLambda and tfd.Beta using the stated TensorFlow and TensorFlow Probability versions, then compare it with the Normal case. Start by tracing tfd.Beta and DistributionLambda around the custom-gradient warning and NaN training metrics. Done means the warning and NaN cause are identified and covered by an appropriate regression test.