Cannot compute gradient of Weibull parameters on second iteration

I am trying to learn the parameters of a Weibull distribution, but after the first iteration the gradients turn None:


import tensorflow as tf
from tensorflow_probability import distributions as tfd

W = tfd.Weibull(concentration=1.5, scale=2.5)
data = W.sample(10_000)

concentration = tf.Variable(initial_value=1.0, dtype=tf.float32, name="W_c")
scale = tf.Variable(initial_value=3.0, dtype=tf.float32, name="W_s")

W_hat = tfd.Weibull(concentration=concentration, scale=scale)

for epoch in range(3):
    with tf.GradientTape() as tape:
        tape.watch(W_hat.trainable_variables)
        nll = -W_hat.log_prob(data)
    grads = tape.gradient(nll, W_hat.trainable_variables)

    print(grads)  # , W_hat.trainable_variables)

Output:

(<tf.Tensor: shape=(), dtype=float32, numpy=585.2751>, <tf.Tensor: shape=(), dtype=float32, numpy=797.33496>)
(None, None)
(None, None)

This does not happen to e.g. the Normal distribution. Please note no optimiser is applied here it's just the computation of the gradient of the nll w.r.t. to the parameters done twice. The versions I'm using:

python 3.8.13 (default, Sep  7 2022, 19:12:31)  [Clang 12.0.5 (clang-1205.0.22.11)]
tensorflow 2.13.0
tensorflow_probability 0.21.0

I reproduced this on both a M1 Mac and an Intel Xeon x86_64

tensorflow / probability

Cannot compute gradient of Weibull parameters on second iteration #1772