Convert to Int array to Binary

[lockwo]

I am struggling to think of a way to convert an array (or matrix) of integers to a matrix of binary values of a certain size that is both jit-able and grad-able. E.g. convert([1, 2, 3], 5) -> [[0, 0, 0, 0, 1], [0, 0, 0, 1, 0], [0, 0, 0, 1, 1]]. I have looked on the online and on the discussion forum and I didn't see anyone who encountered this before. My current implementation is:

def to_binary(data, binsize):
    return jnp.flip(
        data[:, None] & (1 << jnp.arange(binsize)) > 0, axis=-1
    ).astype("int8")

Which computes what I want (and works with jit) but it doesn't seem to work with grad. There is a binary_repr function in numpy but not in jax. There is also an unpack bits function, but that isn't quite the same as this problem. I am curious if anyone has encountered this before, or has any suggestions. Thanks.

[jakevdp]

You can probably implement what you want using the jnp.unpackbits function:

import jax.numpy as jnp
x = jnp.array([1, 2, 3])
bits = jnp.unpackbits(x.astype('uint8')).reshape(len(x), 8)[:, 3:]
print(bits)

[[0 0 0 0 1]
 [0 0 0 1 0]
 [0 0 0 1 1]]

Note that this will not work with grad, because the inputs and outputs are integers. I'm not sure it's really meaningful to talk about the gradient of an operation like this (gradients involve taking limits of infinitessimal changes, which rules-out integers) If there's a particular application you have in mind, it might help to say more about it.

[lockwo]

Thanks for the response! I can talk a little more about what I am trying to do, to explain why I want the gradient. This is basically part of a quantum circuit simulator. So given a statevector, I want to sample from it (and then compute expectation values of operators). A statevector is a 2^N, where N is an integer number of qubits, array of complex values that you square to get a probability distribution. Each qubit can be in a |0> or |1> state, so this is probability distribution over all (2^N) states. E.g. a 2 qubit system has statevector (sv): sv[0] |00> + sv[1] |01> + sv[2] |10> + sv[3] |11>. On real devices you sample from each qubit and get a binary result (so a collection of bit strings), so you might get [[0, 1], [0, 0], [0, 0], [1, 1], etc.] for a thousand samples or so. You can compute the expected value of that bit by doing sum(bit column)/num samples. In my implementation in jax, I sample the (squared) statevector using the random.categorical function (to get random bit strings from the discrete distribution but in integer form, if there is a way to sample bit strings randomly this would also be helpful). I then convert to binary to get the qubit wise (e.g. N expected values) expected value just by doing a jnp.sum operation as described above (then an operator can then act on this, i.e. do a matmul, to get the final result). I think everything up to (and after) the binarization is differentiable, and while methods such as finite difference can be used, I've thought maybe jax could autograd through it (based on reading about grad and random functions here: #3702). That's a small overview of the larger goal I have in mind (hopefully it's clear, not sure on everyone's QC background, let me know if anything I can explain better).

I agree, gradient of integer functions aren't super meaningful, but this binarization function is used to generate samples (basically, this is a small part to taking the gradient of a noisy/stochastic function) of floating point values which the gradient does matter for. TLDR of the problem is:

Differentaible function generates a discrete probability distribution -> sample from it using categorical (which I thought was differentiable per #3702) -> convert to binary -> average over binary columns -> do a matmul -> floating point loss value.

Everything seems to be working with grad in this workflow except convert to binary.

[jakevdp]

Thanks for the details - that makes sense. In that case, I think you'll probably have to define your own function via a custom_jvp in order to integrate through the binarization. I don't have any great ideas on how to compute the gradients of such a function, but I'd probably start by writing down symbolically what you think the gradient should be, and try to implement the result via custom_jvp.

[NoraAl]

@jakevdp Is there an implementation for numbers more than 8 bits, or a way to indicate an overflow?

[jakevdp]

Sure, you can view larger integers as uint8 and then unpack the bits. For example:

x = jnp.arange(4, dtype='int32')
print(jnp.unpackbits(x.view('uint8'), bitorder='little').reshape(4, 32))
# [[0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]
#  [1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]
#  [0 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]
#  [1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]]

[NoraAl]

My target array was fixed and small, so it didn't show the whole representation. I should've tried it on bigger targets. Very helpful, thanks.