categorical.py 文件源码

def _log_prob(self, k):
    k = ops.convert_to_tensor(k, name="k")
    logits = self.logits * array_ops.ones_like(
        array_ops.expand_dims(k, -1),
        dtype=self.logits.dtype)
    shape = array_ops.slice(array_ops.shape(logits), [0],
                            [array_ops.rank(logits) - 1])
    k *= array_ops.ones(shape, dtype=k.dtype)
    k.set_shape(tensor_shape.TensorShape(logits.get_shape()[:-1]))
    return -nn_ops.sparse_softmax_cross_entropy_with_logits(logits, k)