Python relaxedonehot分类的梯度似乎为零

我试图预先训练一个seq-to-seq编码器-解码器模型，除了使用不同的数据集外，该模型非常接近此处列出的结构（）。我计划将此模型作为GAN模型中的生成器，这意味着我需要以可微的方式从模型中采样文本（即，我不能使用任何argmax函数）。我注意到Gumbel采样是一个潜在的解决方案，Tensorflow在tfp.distributions.RelaxedoneHotCategory类中实现了Gumbel采样。然而，当我在模型中实现该类时，培训似乎不再收敛。解码器类编写如下：

class Decoder(tf.keras.Model):
    def __init__(self, vocab_size, embedding_dim, dec_units, batch_sz, temp):
        super(Decoder, self).__init__()
        self.batch_sz = batch_sz
        self.dec_units = dec_units
        self.embedding = tf.keras.layers.Embedding(vocab_size, embedding_dim)
        self.gru = tf.keras.layers.GRU(self.dec_units,
                                       return_sequences=True,
                                       return_state=True,
                                       recurrent_initializer='glorot_uniform')
        self.fc = tf.keras.layers.Dense(vocab_size)
        self.gumbel = tfp.layers.DistributionLambda(lambda logits: tfp.distributions.RelaxedOneHotCategorical(logits=logits, temperature=temp), convert_to_tensor_fn=lambda s: s.sample())
        self.temp = temp
        self.attention = BahdanauAttention(self.dec_units)
        
    def __call__(self, x, hidden, enc_output):
        context_vector, attention_weights = self.attention(hidden, enc_output)
        x = self.embedding(x)
        x = tf.concat([tf.expand_dims(context_vector, 1), x], axis=-1)
        output, state = self.gru(x)
        output = tf.reshape(output, (-1, output.shape[2]))
        x = self.fc(output)
        x = self.gumbel(x)
        x_hard = tf.stop_gradient(tf.cast(tf.equal(x,tf.reduce_max(x,1,keepdims=True)),x.dtype))
        x = tf.stop_gradient(x_hard - x) + x

        return x, state, attention_weights

调用方法的最后几行使用argmax为前向传递生成一个热输出，但是由于调用了

stop\u gradient（x\u hard-x）

，Tensorflow应该忽略后向传递中的argmax，并计算到relaxedonehotCategory的梯度w.r.t。然而，我仍然遇到了培训衔接问题

x_hard = tf.stop_gradient(tf.cast(tf.equal(x,tf.reduce_max(x,1,keepdims=True)),x.dtype))
x = tf.stop_gradient(x_hard - x) + x

台词是不恰当的

即使我尝试下面这样的简单小示例，我也看到返回的渐变仍然是0

@tf.function
def test(a):
    b = tf.nn.sigmoid(a)
    c = tfp.distributions.RelaxedOneHotCategorical(logits=b, temperature=.1).sample()
    return tf.gradients(c, [a, b])

a = tf.Variable(np.array([.1, .2, .7]))
test(a)

有人知道如何解决这个问题吗？作为参考，我使用的是Tensorflow 2.3和Tensorflow概率0.11.1