model.py 文件源码

def multinomial_3d(x):
    """Samples from a multinomial distribution from 3D Tensor.
    Args:
        x: Tensor with shape (batch_size, timesteps, classes)
    Returns:
        Tensor with shape (batch_size, timesteps), sampled from `classes`.
    """
    a, b = tf.shape(x)[0], tf.shape(x)[1]
    x = tf.reshape(x, (a * b, -1))
    m = tf.multinomial(x, 1)
    return tf.reshape(m, (a, b))