INNER CODE UNIT · Python

log_prob

lucidrains/vector-quantize-pytorch · vector_quantize_pytorch/binary_mapper.py:89

    def log_prob(
        self,
        logits,
        *,
        indices = None,
        one_hot = None,
        sum_bits = True
    ):
        assert exists(indices) ^ exists(one_hot), 'either indices or one_hot must be provided'

        if exists(one_hot):
            indices = one_hot.argmax(dim=-1)

        # allow for any number of leading dimensions

        logits, inverse_pack_lead_dims = pack_with_inverse(logits, '* bits')
        indices, _ = pack_with_inverse(indices, '*')

View source record →