I’m afraid I’m not sure what the issue is here. Keras does some kind of tracing that I don’t understand, and it seems not to like one of the functions we’re using there. Pasting the full stack trace below. It seems like some crucial information is being lost at the point where the real error is detected. Maybe an exception is being caught and re-raised, losing some context? Maybe someone from Keras can weigh in? I don’t think the TFP code is doing anything unusual here.
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-17-c3244dc752a3> in <module>()
     18 
     19 Y = Input(shape=(1,))
---> 20 lossF =  betabin_mdn_cost(a,b,Y)
     21 model3 = tf.keras.Model(inputs=[InputLayer, Y], outputs=[a,b])
     22 model3.add_loss(lossF)

<ipython-input-17-c3244dc752a3> in betabin_mdn_cost(a_, b_, y, total_count)
      6                                          concentration0=a_,
      7                                          concentration1=b_) 
----> 8    return tf.reduce_mean(-dist.log_prob(y))
      9 
     10 epochs = 500

google3/third_party/py/tensorflow_probability/python/distributions/distribution.py in log_prob(self, value, name, **kwargs)
   1282         values of type `self.dtype`.
   1283     """
-> 1284     return self._call_log_prob(value, name, **kwargs)
   1285 
   1286   def _call_prob(self, value, name, **kwargs):

google3/third_party/py/tensorflow_probability/python/distributions/distribution.py in _call_log_prob(self, value, name, **kwargs)
   1264     with self._name_and_control_scope(name, value, kwargs):
   1265       if hasattr(self, '_log_prob'):
-> 1266         return self._log_prob(value, **kwargs)
   1267       if hasattr(self, '_prob'):
   1268         return tf.math.log(self._prob(value, **kwargs))

google3/third_party/py/tensorflow_probability/python/internal/distribution_util.py in _fn(*args, **kwargs)
   1364     @functools.wraps(fn)
   1365     def _fn(*args, **kwargs):
-> 1366       return fn(*args, **kwargs)
   1367 
   1368     if _fn.__doc__ is None:

google3/third_party/py/tensorflow_probability/python/distributions/beta_binomial.py in _log_prob(self, counts)
    271     return (_log_combinations(n, counts)
    272             + tfp_math.lbeta(c1 + counts, (n - counts) + c0)
--> 273             - tfp_math.lbeta(c1, c0))
    274 
    275   @distribution_util.AppendDocstring(_beta_binomial_sample_note)

google3/third_party/py/tensorflow_probability/python/distributions/beta_binomial.py in _log_combinations(n, k)
    329 def _log_combinations(n, k):
    330   """Computes log(Gamma(n+1) / (Gamma(k+1) * Gamma(n-k+1))."""
--> 331   return -tfp_math.lbeta(k + 1, n - k + 1) - tf.math.log(n + 1)

google3/third_party/tensorflow/python/util/traceback_utils.py in error_handler(*args, **kwargs)
    151     except Exception as e:
    152       filtered_tb = _process_traceback_frames(e.__traceback__)
--> 153       raise e.with_traceback(filtered_tb) from None
    154     finally:
    155       del filtered_tb

google3/third_party/py/keras/engine/keras_tensor.py in __array__(self)
    243   def __array__(self):
    244     raise TypeError(
--> 245         'Cannot convert a symbolic Keras input/output to a numpy array. '
    246         'This error may indicate that you\'re trying to pass a symbolic value '
    247         'to a NumPy call, which is not supported. Or, '

TypeError: Cannot convert a symbolic Keras input/output to a numpy array. This error may indicate that you're trying to pass a symbolic value to a NumPy call, which is not supported. Or, you may be trying to pass Keras symbolic inputs/outputs to a TF API that does not register dispatching, preventing Keras from automatically converting the API call to a lambda layer in the Functional Model.