How to use 1-bit booleans like C++'s bitset

[imoneoi]

Can I use JAX to compress 32 boolean variables into one uint32 variable, like C++'s bitset functionality? I want to implement architectures like binary neural networks in JAX.

For example, I have a boolean array of length 256 that I want to compress into a uint32 array of length 8.

[jakevdp]

XLA doesn't have any support for 1-bit variables, but recent versions of JAX do have some experimental suppot for 4-bit integers via the int4 dtype that works on supported hardware:

In [1]: import jax.numpy as jnp

In [2]: x = jnp.arange(4, dtype='int4')

In [3]: print(x)
[0 1 2 3]

The only way I know of to use bit arrays is via the jnp.packbits functionality, which looks something like this:

In [4]: x = jnp.array([1, 0, 1, 1, 0, 0, 1, 0], dtype='uint8')

In [5]: bits = jnp.packbits(x)

In [6]: print(bits)
[178]

In [7]: print(jnp.unpackbits(bits))
[1 0 1 1 0 0 1 0]

[imoneoi]

Thanks! I think we might currently use 1-bit datatypes and matrix multiplication like in the example below (cf. packbits). What do you recommend?

def bitset_from_bool(x, bitset_dtype):
    """Create bitset from boolean (on last dimension)."""
    assert (x.shape[-1] % jnp.iinfo(bitset_dtype).bits) == 0, "Array length must be multiple of bits of bitset dtype."

    # To destination dtype
    x = x.view(bitset_dtype)

    # Bitwise or
    bits = jnp.arange(8, jnp.uint8)
    x    = x.reshape(x.shape[:-1] + (x.shape[-1] // 8, 8))
    x    = jax.lax.reduce(x << bits, 0, jax.lax.bitwise_or, (-1, ))
    return x


def bitset_matvec_mul(A, b, out_dtype):
    """Bitset matrix-vector multiplication Y = Ab (on last dimension).
    A: ... x N x (M/Q) matrix
    B: ... x (M/Q)     vector 
    Y: ... x N         vector"""
    assert A.dtype == b.dtype, "A and b must have the same dtype."
    assert A.shape[-1] == b.shape[-1], "The last dimension of A and b must be of same size."

    # Broadcasted bitwise and + popcount
    x = jax.lax.bitwise_and(A, jnp.expand_dims(b, -2))
    # Popcount
    x = jax.lax.population_count(x)
    # Accumulate (reduce sum)
    return jax.lax.reduce(x.astype(out_dtype), 0, jax.lax.add, (-1, ))