INNER CODE UNIT · Python

BinaryMapper

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

class BinaryMapper(Module):
    def __init__(
        self,
        bits = 1,
        kl_loss_threshold = NAT, # 1 bit
        deterministic_on_eval = False
    ):
        super().__init__()

        self.bits = bits
        self.num_codes = 2 ** bits

        power_two = 2 ** arange(bits)
        codes = (arange(self.num_codes)[:, None].bitwise_and(power_two) != 0).byte().bool()

        self.register_buffer('power_two', power_two, persistent = False)
        self.register_buffer('codes', codes, persistent = False)

View source record →