Jasp-compatible Montgomery arithmetics

[cjn10]

Adds jasp-compatible (and non-jasp-compatible) functions for montgomery arithmetics. Additionally, removes the length-check in the qq_gidney_adder.

[positr0nium]

First of all: Massive work!! This is amazing. It's clear that you put a lot of work in it and this kind of feature underlines once more how we are truly pushing the state of the art in quantum compilation.

Since this is a big commit, I have a couple of points, this however doesn't at all takes anything away from the fact how cool this is :)

• Moving forward, we should start testing arithmetic with boolean simulation which allows us to test much more cases (not just one).
• Test montgomery_jasp_qq seems to not properly uncompute (terminal message during simulation).
• get_size function in jasp_montgomery is already available in the global scope as jlen
• Python Big Integer Import fails for >64 bits because the dynamic loop needs to convert to jnp.in64 first. Using a simple Python loop would suffice as a fix.

from qrisp import jaspify
from qrisp.alg_primitives.arithmetic.jasp_arithmetic import BigInteger

b_int = 532_323_451_122_411

@jaspify
def test():
    b = BigInteger.from_float(b_int, size=16)
    return (b)()

assert ((float(b_int)) == test())

• BigInteger.__add__ could be more efficient via bitshift instead of floordivision
• BigInteger.__sub__ might be more efficient via a - b = NOT ((NOT a + b) mod N)

n = 5
N = 1<<n
N_prime = N - 1

a = 10
b = 4

# apply NOT on a
a = a ^ N_prime

# addition modulo N
c = (a + b) & N_prime

# apply NOT on c
c = c ^ N_prime

# Yields a - b
print(c)

• Type checking in QuantumModolus.__imul__ doesn't seem to be working for traced integers and big integers.
• T-gate count for controlled in-place multiplication is much better than in static mode. Not a problem of course but interesting 🤔

n = 5
X = 29
y = 21
N = 31
n = 5

@count_ops(meas_behavior = "1")
def test_cq_ft():
    qy = QuantumFloat(n)
    qy[:] = y
    m = compute_aux_radix_exponent(N, qy.size)
    ctrl = QuantumBool()
    
    with control(ctrl):
        cq_montgomery_multiply_inplace(X, qy, N, m, gidney_adder)
    return measure(qy)

ops = test_cq_ft()
print(ops)

qy = QuantumModulus(N, inpl_adder = gidney_adder)
qy[:] = y
m = compute_aux_radix_exponent(N, qy.size)

ctrl = QuantumBool()
    
with control(ctrl):
    qy *= X

qc = qy.qs.compile()

print(qc.transpile().count_ops())

[cjn10]

Thank you very much for the feedback. I have pushed some changes that address the points you have made.

• Moving forward, we should start testing arithmetic with boolean simulation which allows us to test much more cases (not just one).

I am working on that.

• Test montgomery_jasp_qq seems to not properly uncompute (terminal message during simulation).

This appears to be a bug in the controlled version of the jasp_qq_gidney_adder, see #263. The uncomputation is successful for the fourier adder.

• get_size function in jasp_montgomery is already available in the global scope as jlen

Fixed.

• Python Big Integer Import fails for >64 bits because the dynamic loop needs to convert to jnp.in64 first. Using a simple Python loop would suffice as a fix.

When executing the privided code, no error is thrown for me. However, I agree that the way to create a BigInteger at the moment is not optimal. I will look into this and try to come up with a nicer solution (both internally and on the user side).

• BigInteger.__add__ could be more efficient via bitshift instead of floordivision

Good catch.

• BigInteger.__sub__ might be more efficient via a - b = NOT ((NOT a + b) mod N)

I have implemented this method and the required bit operations. However, I think the added overhead by calling the bit operations is larger than that of the direct implementation.

• Type checking in QuantumModolus.__imul__ doesn't seem to be working for traced integers and big integers.

Fixed.

• T-gate count for controlled in-place multiplication is much better than in static mode. Not a problem of course but interesting

I agree that this is interesting. This might be caused by different default adders. I will investigate further!

[positr0nium]

Regarding the big_int initialization: You're right, the presented example doesn't reproduce the error. No idea why I posted that. This one however gives me an error:

b_int = 532_323_451_122_411_593_564_574_574

[cjn10]

Regarding the big_int initialization: You're right, the presented example doesn't reproduce the error. No idea why I posted that. This one however gives me an error:

b_int = 532_323_451_122_411_593_564_574_574

Yes, this is indeed a problem. The jax loop inside the function requires a dynamic jax integer which are limited in size (unlike Python integers). In this case, however, we can use a Python loop to create the BigInteger using the function from_int_python. I will make this more clear and obvious in a future version.

[renezander90]

import numpy as np
import jax
from qrisp import terminal_sampling, QuantumModulus, QuantumFloat, jrange, control, jasp_fourier_adder, QFT, h, x, fourier_adder, gidney_adder, jasp_cq_gidney_adder, BigInteger, remaud_adder

def find_order(a, N):
    qg = QuantumModulus(N, inpl_adder=remaud_adder)
    qg[:] = 1
    qpe_res = QuantumFloat(2*qg.size + 1, exponent=-(2*qg.size + 1))
    h(qpe_res)
    for i in range(len(qpe_res)):
        with control(qpe_res[i]):
            qg *= a
            a = (a*a) % N
    QFT(qpe_res, inv=True)
    return qpe_res.get_measurement()

dict_norm = find_order(4, 13)

Yields: TypeError: remaud_adder() missing 1 required positional argument: 'z'

Not sure if this is a problem, with the remand adder, or caused by using this adder within the jasp_montgomery arithmetic. (The change was made to always use the jasp Montgomery arithmetic for quantum-classical.)

[renezander90]

An example like this, demonstrating how BigInteger can be used with QuantumModulus, would be helpful in the Scalable Integer Type documentation.

from qrisp import *

@boolean_simulation
def main():

    N = BigInteger.create_static(340282366762482138434845932244680310783)
    a = BigInteger.create_static(131313)

    qm = QuantumModulus(N)
    qm[:] = 1

    for i in jrange(10):
        qm *= a

    return measure(qm)()

main()

However, even running a smaller example, like this:

from qrisp import *

@boolean_simulation
def main():

    N = BigInteger.create_static(340282366762482138434845932244680310783)

    qm = QuantumModulus(N)
    qm[:] = 1
    qm *= 21

    return measure(qm)()

main()

Yields the correct result, but:

WARNING: Faulty uncomputation found during simulation.
WARNING: Faulty uncomputation found during simulation.
WARNING: Faulty uncomputation found during simulation.
WARNING: Faulty uncomputation found during simulation.
WARNING: Faulty uncomputation found during simulation.
WARNING: Faulty uncomputation found during simulation.

Array(21., dtype=float64)

[positr0nium]

import numpy as np
import jax
from qrisp import terminal_sampling, QuantumModulus, QuantumFloat, jrange, control, jasp_fourier_adder, QFT, h, x, fourier_adder, gidney_adder, jasp_cq_gidney_adder, BigInteger, remaud_adder

def find_order(a, N):
    qg = QuantumModulus(N, inpl_adder=remaud_adder)
    qg[:] = 1
    qpe_res = QuantumFloat(2*qg.size + 1, exponent=-(2*qg.size + 1))
    h(qpe_res)
    for i in range(len(qpe_res)):
        with control(qpe_res[i]):
            qg *= a
            a = (a*a) % N
    QFT(qpe_res, inv=True)
    return qpe_res.get_measurement()

dict_norm = find_order(4, 13)

Yields: TypeError: remaud_adder() missing 1 required positional argument: 'z'

Not sure if this is a problem, with the remand adder, or caused by using this adder within the jasp_montgomery arithmetic. (The change was made to always use the jasp Montgomery arithmetic for quantum-classical.)

@renezander90 The Remaud adder still needs some more work to be used here. What is missing in particular is a classical-quantum addition. Considering it uses the input qubits as "dirty workspace" a lot, I don't think a more optimal version than this can be found:

from qrisp import *

qf = QuantumFloat(5)
qf[:] = 4
def cl_remaud_adder(c, qf):
    temp = qf.duplicate()
    z = [] # This seems to work??
    
    with conjugate(int_encoder)(temp, c):
        remaud_adder(temp, qf, z)
        
    temp.delete()
    
cl_remaud_adder(5, qf)

print(qf)

Furthermore we need to figure out how to deal with the carry value z. In the snippet it seems to work but not sure about edge cases. (Tagging @diehoq in case he is interested).

[renezander90]

classical-quantum matrix multiplication for QuantumArray with type QuantumModulus does not work properly (are there any tests?):

import jax.numpy as jnp
import numpy as np

def main():

    a = np.array([[1,2],[3,4]])

    b = QuantumArray(QuantumModulus(13),shape=(2,2))
    b[:] = np.array([[1,2],[3,4]])
    
    c = a @ b

    return c.get_measurement()

main()

Yields:

ValueError: cannot convert float NaN to integer

from qrisp import *
import jax.numpy as jnp
import numpy as np

@boolean_simulation
def main():

    a = np.array([[1,2],[3,4]])

    b = QuantumArray(QuantumModulus(13),shape=(2,2))
    b[:] = np.array([[1,2],[3,4]])
    
    c = a @ b

    return measure(c)

main()

Yields:

KeyError: <QuantumModulus 'qf_0'>