Error when using ZeroInflatedNegativeBinomial inside DistributionLambda

[Guilherme Namen Pimenta]

Hi, I am trying to use the ZeroInflatedNegativeBinomial distribution inside the DistributionLambda layer using the flowing code:

x = layers.Dense(DENSE_SIZE)(all_features)
x = layers.BatchNormalization()(x)
x = layers.Activation(activation=ACTIVATION)(x)
x = layers.Dropout(DROP_OUT)(x)
x = layers.Dense(3)(x)
output = tfpl.DistributionLambda(
    make_distribution_fn=lambda t:  tfd.ZeroInflatedNegativeBinomial(total_count=tf.exp(t[..., 0]), logits=t[..., 1], inflated_loc_logits=t[..., 2]),
    convert_to_tensor_fn=tfd.Distribution.mean)(x)

But, when I compile the model it rises the following error:

TypeError Traceback (most recent call last) Cell In [19], line 17 15 x = layers.Dropout(DROP_OUT)(x) 16 x = layers.Dense(3)(x) ---> 17 output = tfpl.DistributionLambda( 18 make_distribution_fn=lambda t: tfd.ZeroInflatedNegativeBinomial(total_count=tf.exp(t[..., 0]), logits=t[..., 1], inflated_loc_logits=t[..., 2]), 19 convert_to_tensor_fn=tfd.Distribution.mean)(x) 21 model = keras.Model(all_inputs, output) 22 model.compile(OPTIMIZER, LOSS, metrics=METRICS) File ~\Python\env\lib\site-packages\tensorflow_probability\python\layers\distribution_layer.py:220, in DistributionLambda.__call__(self, inputs, *args, **kwargs) 218 def __call__(self, inputs, *args, **kwargs): 219 self._enter_dunder_call = True --> 220 distribution, _ = super(DistributionLambda, self).__call__( 221 inputs, *args, **kwargs) 222 self._enter_dunder_call = False 223 return distribution File ~\Python\env\lib\site-packages\keras\utils\traceback_utils.py:70, in filter_traceback.<locals>.error_handler(*args, **kwargs) 67 filtered_tb = _process_traceback_frames(e.__traceback__) 68 # To get the full stack trace, call: 69 # `tf.debugging.disable_traceback_filtering()` ---> 70 raise e.with_traceback(filtered_tb) from None 71 finally: 72 del filtered_tb File ~\Python\env\lib\site-packages\tensorflow_probability\python\layers\distribution_layer.py:226, in DistributionLambda.call(self, inputs, *args, **kwargs) 225 def call(self, inputs, *args, **kwargs): --> 226 distribution, value = super(DistributionLambda, self).call( 227 inputs, *args, **kwargs) 228 # We always save the most recently built distribution variables for tracking 229 # purposes. 230 self._most_recently_built_distribution_vars = distribution.variables File ~\Python\env\lib\site-packages\tensorflow_probability\python\layers\distribution_layer.py:180, in DistributionLambda.__init__.<locals>._fn(*fargs, **fkwargs) 171 distribution = dtc._TensorCoercible( # pylint: disable=protected-access 172 distribution=d, 173 convert_to_tensor_fn=maybe_composite_convert_to_tensor_fn) 175 # Calling `distrbution._value()` is equivalent to: 176 # from tensorflow.python.framework import ops 177 # value = ops.convert_to_tensor_or_composite(distribution) 178 # We'd prefer to call ops.convert_to_tensor_or_composite but do not, 179 # favoring our own non-public API over TF's. --> 180 value = distribution._value() # pylint: disable=protected-access 182 # TODO(b/126056144): Remove silent handle once we identify how/why Keras 183 # is losing the distribution handle for activity_regularizer. 184 value._tfp_distribution = distribution # pylint: disable=protected-access File ~\Python\env\lib\site-packages\tensorflow_probability\python\layers\internal\distribution_tensor_coercible.py:215, in _TensorCoercible._value(self, dtype, name, as_ref) 206 raise NotImplementedError( 207 'Failed to convert object of type {} to Tensor. Contents: {}. ' 208 'Call `distribution.set_tensor_conversion(lambda self: ...)` to ' (...) 211 ' results in `tf.convert_to_tensor(x)` being identical to ' 212 '`x.mean()`.'.format(type(self), self)) 213 with self._name_and_control_scope('value'): 214 self._concrete_value = ( --> 215 self._convert_to_tensor_fn(self.tensor_distribution) 216 if callable(self._convert_to_tensor_fn) 217 else self._convert_to_tensor_fn) 218 if (not tf.is_tensor(self._concrete_value) and 219 not isinstance(self._concrete_value, 220 composite_tensor.CompositeTensor)): 221 self._concrete_value = nest_util.convert_to_nested_tensor( # pylint: disable=protected-access 222 self._concrete_value, 223 name=name or 'concrete_value', 224 dtype=dtype, 225 dtype_hint=self.tensor_distribution.dtype) File ~\Python\env\lib\site-packages\tensorflow_probability\python\distributions\distribution.py:1532, in Distribution.mean(self, name, **kwargs) 1530 """Mean.""" 1531 with self._name_and_control_scope(name): -> 1532 return self._mean(**kwargs) File ~\Python\env\lib\site-packages\tensorflow_probability\python\distributions\mixture.py:251, in _Mixture._mean(self) 250 def _mean(self): --> 251 distribution_means = [d.mean() for d in self.components] 252 cat_probs = self._cat_probs(log_probs=False) 253 cat_probs = [self._expand_to_event_rank(c_p) for c_p in cat_probs] File ~\Python\env\lib\site-packages\tensorflow_probability\python\distributions\mixture.py:251, in <listcomp>(.0) 250 def _mean(self): --> 251 distribution_means = [d.mean() for d in self.components] 252 cat_probs = self._cat_probs(log_probs=False) 253 cat_probs = [self._expand_to_event_rank(c_p) for c_p in cat_probs] File ~\Python\env\lib\site-packages\tensorflow_probability\python\distributions\distribution.py:1532, in Distribution.mean(self, name, **kwargs) 1530 """Mean.""" 1531 with self._name_and_control_scope(name): -> 1532 return self._mean(**kwargs) File ~\Python\env\lib\site-packages\tensorflow_probability\python\distributions\deterministic.py:138, in _BaseDeterministic._mean(self) 137 def _mean(self): --> 138 loc = tf.convert_to_tensor(self.loc) 139 return tf.broadcast_to( 140 loc, 141 ps.concat([self._batch_shape_tensor(loc=loc), 142 self._event_shape_tensor(loc=loc)], 143 axis=0)) File ~\Python\env\lib\site-packages\tensorflow_probability\python\util\deferred_tensor.py:83, in _tensorize(d, dtype, name, as_ref) 81 def _tensorize(d, dtype=None, name=None, as_ref=False): 82 """Tensor conversion function presuming `hasattr(d, '_value')`.""" ---> 83 return d._value(dtype, name, as_ref) File ~\Python\env\lib\site-packages\tensorflow_probability\python\util\deferred_tensor.py:367, in DeferredTensor._value(self, dtype, name, as_ref) 366 def _value(self, dtype=None, name=None, as_ref=False): --> 367 y = self.transform_fn(self.pretransformed_input) # pylint: disable=not-callable 368 if dtype_util.base_dtype(y.dtype) != self.dtype: 369 raise TypeError( 370 'Actual dtype ({}) does not match deferred dtype ({}).'.format( 371 dtype_util.name(dtype_util.base_dtype(y.dtype)), 372 dtype_util.name(self.dtype))) File ~\Python\env\lib\site-packages\tensorflow_probability\python\distributions\inflated.py:131, in _Inflated.__init__.<locals>.<lambda>(x) 119 categorical_dist = categorical.Categorical( 120 probs=cat_probs, 121 validate_args=validate_args, 122 allow_nan_stats=allow_nan_stats) 123 probs_or_logits = self._inflated_loc_probs 125 super(_Inflated, self).__init__( 126 cat=categorical_dist, 127 components=[ 128 deterministic.Deterministic( 129 DeferredTensor( 130 probs_or_logits, --> 131 lambda x: tf.constant( # pylint: disable=g-long-lambda 132 inflated_loc, dtype=distribution.dtype, 133 shape=probs_or_logits.shape)), 134 validate_args=validate_args, 135 allow_nan_stats=allow_nan_stats), 136 distribution 137 ], 138 validate_args=validate_args, 139 allow_nan_stats=allow_nan_stats, 140 name=name) 141 self._parameters = parameters File <__array_function__ internals>:180, in prod(*args, **kwargs) File ~\Python\env\lib\site-packages\numpy\core\fromnumeric.py:3045, in prod(a, axis, dtype, out, keepdims, initial, where) 2927 @array_function_dispatch(_prod_dispatcher) 2928 def prod(a, axis=None, dtype=None, out=None, keepdims=np._NoValue, 2929 initial=np._NoValue, where=np._NoValue): 2930 """ 2931 Return the product of array elements over a given axis. 2932 (...) 3043 10 3044 """ -> 3045 return _wrapreduction(a, np.multiply, 'prod', axis, dtype, out, 3046 keepdims=keepdims, initial=initial, where=where) File ~\Python\env\lib\site-packages\numpy\core\fromnumeric.py:86, in _wrapreduction(obj, ufunc, method, axis, dtype, out, **kwargs) 83 else: 84 return reduction(axis=axis, out=out, **passkwargs) ---> 86 return ufunc.reduce(obj, axis, dtype, out, **passkwargs) TypeError: Exception encountered when calling layer "distribution_lambda_4" (type DistributionLambda). int() argument must be a string, a bytes-like object or a real number, not 'NoneType' Call arguments received by layer "distribution_lambda_4" (type DistributionLambda): • inputs=tf.Tensor(shape=(None, 3), dtype=float32) • args=<class 'inspect._empty'> • kwargs={'training': 'False'}

[Colin Carroll]

Can you share more about 

1. What versions of the libraries are you using? (tfp.__version__, tf.__version__)

2. A minimal example, maybe in a colab, that reproduces this error? 

For what it is worth, I ran, after !pip install -Uq tfp-nightly tf-nightly, the following successfully:

import numpy as np

import tensorflow as tf

import tensorflow_probability as tfp

from keras import layers

tfpl = tfp.layers

tfd = tfp.distributions

print(tf.__version__, tfp.__version__)

DENSE_SIZE = 32

DROP_OUT = 0.5

ACTIVATION = 'relu'

all_features = tf.convert_to_tensor(np.random.rand(100, 16), tf.float32)

x = layers.Dense(DENSE_SIZE)(all_features)

x = layers.BatchNormalization()(x)

x = layers.Activation(activation=ACTIVATION)(x)

x = layers.Dropout(DROP_OUT)(x)

x = layers.Dense(3)(x)

output = tfpl.DistributionLambda(

make_distribution_fn=lambda t: tfd.ZeroInflatedNegativeBinomial(

total_count=tf.exp(t[..., 0]),

logits=t[..., 1],

inflated_loc_logits=t[..., 2]),

convert_to_tensor_fn=tfd.Distribution.mean)(x)

[Guilherme Namen Pimenta]

Hi, thanks a lot! Here I am using:

      TF version: 2.10.0

      TFP version: 0.18.0

The example was attached.