Looking at the stacktrace @Christopher_Suter shared: I think some traceback filtering got added to TF recently that is removing some callframes that would maybe also be useful here. Taking a look though, it looks like that tfp_math.lbeta method in the stacktrace actually calls a tf.function:

  

      github.com
  

  
    tensorflow/probability/blob/main/tensorflow_probability/python/math/special.py#L1322



      
      
          @tfp_custom_gradient.custom_gradient(
              vjp_fwd=_lbeta_fwd,
              vjp_bwd=_lbeta_bwd,
              jvp_fn=_lbeta_jvp)
          def _lbeta_custom_gradient(x, y):
            """Computes log(Beta(x, y)) with correct custom gradient."""
            return _lbeta_naive_gradient(x, y)
          
          

          
@tf.function(autograph=False)
          def lbeta(x, y, name=None):
            """Returns log(Beta(x, y)).
          
          
  This is semantically equal to
              lgamma(x) + lgamma(y) - lgamma(x + y)
            but the method is more accurate for arguments above 8.
          
          
  The reason for accuracy loss in the naive computation is catastrophic
            cancellation between the lgammas.  This method avoids the numeric cancellation
            by explicitly decomposing lgamma into the Stirling approximation and an
            explicit `log_gamma_correction`, and cancelling the large terms from the
      
    
  


  

  
    
    
  

  


It is not supported to call tf.function directly on symbolic intermediate values during Functional model construction (hence the KerasTensors error message). What you can do is put the tf.function call / offending piece of code in a Keras lambda layer or a custom Keras layer’s call method. (In this case that would be putting betabin_mdn_cost(a, b, Y) in the call method of a custom layer that takes three inputs)
We’ll update this error message to mention that as the workaround whenever it gets raised.