title: Tensorflow Probability Regression Setup
body: <p>I'm trying to build upon case 2 in <a href="https://www.tensorflow.org/probability/examples/Probabilistic_Layers_Regression" rel="nofollow noreferrer">https://www.tensorflow.org/probability/examples/Probabilistic_Layers_Regression</a></p>
  <p>Say I have a dataframe, I want set up a model like:</p>
  <p>Y ~ NegativeBinomial(a,b),</p>
  <p>where:</p>
  <p>a ~ kerasNN(X), b ~ kerasNN(X)</p>
  <p>in much the same way they've done for the Normal distribution in case 2.</p>
  <p>However when I do this:</p>
  <pre><code>model = tf.keras.Sequential([
    tf.keras.layers.Dense(1 + 1),
    tfp.layers.DistributionLambda(
        lambda t: tfd.NegativeBinomial(
            total_count=t[..., :1],
            probs=1e-3 + tf.math.softplus(0.01 * t[...,1:])
        )
    ),
  ])
  </code></pre>
  <p>I get loss: nan.</p>
  <p>Does anyone have a template for this sort of problem or know any accessible reading material for setting this up correctly?</p>
  <p>Many thanks.</p>