fixed a bug that prevented proper compilation of balauca mcx in traced mode of static lists of controls

Author: positr0nium

src/qrisp/alg_primitives/mcx_algs/balauca.py

@@ -29,7 +29,7 @@
 from qrisp.alg_primitives.mcx_algs.circuit_library import reduced_maslov_qc, margolus_qc, reduced_margolus_qc
 from qrisp.alg_primitives.mcx_algs.gidney import GidneyLogicalAND
 from qrisp.environments import invert, control, conjugate
-from qrisp.jasp import check_for_tracing_mode, AbstractQubit, qache, jrange, make_tracer
+from qrisp.jasp import check_for_tracing_mode, AbstractQubit, qache, jrange, make_tracer, jlen
 
 # Ancilla supported multi controlled X with logarithmic depth based on
 # https://www.iccs-meeting.org/archive/iccs2022/papers/133530169.pdf
@@ -447,19 +447,28 @@ def extract_boolean_digit(integer, digit):
 
 
 def ctrl_state_conjugator(ctrls, ctrl_state):
-    for i in jrange(ctrls.size):
+    
+    if isinstance(ctrls, list):
+        xrange = range
+    else:
+        xrange = jrange
+    
+    N = jlen(ctrls)
+    
+    for i in xrange(N):
         with control(~extract_boolean_digit(ctrl_state, i)):
             x(ctrls[i])
 
 
 @qache
 def jasp_balauca_mcx(ctrls, target, ctrl_state):
 
+    N = jlen(ctrls)
+    
     from qrisp import mcx
     ctrl_state = jnp.int64(ctrl_state)
-    ctrl_state = cond(ctrl_state == -1, lambda x : x + 2**ctrls.size, lambda x : x, ctrl_state)
+    ctrl_state = cond(ctrl_state == -1, lambda x : x + 2**N, lambda x : x, ctrl_state)
 
-    N = ctrls.size
 
     with conjugate(ctrl_state_conjugator)(ctrls, ctrl_state):
 
@@ -506,23 +515,31 @@ def jasp_balauca_mcp(phi, ctrls, ctrl_state):
 def jasp_balauca_helper(ctrls, balauca_anc):
     from qrisp import mcx
 
-    N = ctrls.size
+    if isinstance(ctrls, list):
+        xrange = range
+        import numpy as jnp
+    else:
+        xrange = jrange
+        import jax.numpy as jnp
+        
+    N = jlen(ctrls)
+    
     n = jnp.int64(jnp.ceil(jnp.log2(N)))
 
-    for i in jrange(N//2):
+    for i in xrange(N//2):
         mcx([ctrls[2*i], ctrls[2*i+1]], balauca_anc[i])
 
-    with control(jnp.bool(N%2)):
+    with control(N%2 != 0):
         cx(ctrls[N-1], balauca_anc[N//2-1+N%2])
 
     n = jnp.int64(jnp.ceil(jnp.log2(N)))
 
     l = make_tracer(0)
     k = N
-    for i in jrange(n-2):
+    for i in xrange(n-2):
         k = jnp.int64(jnp.ceil(k/2))
 
-        for j in jrange(k//2):
+        for j in xrange(k//2):
             mcx([balauca_anc[l+2*j], balauca_anc[l+2*j+1]], 
                 balauca_anc[l+k+j], 
                 method = "gidney")

src/qrisp/environments/control_environment.py

@@ -474,7 +474,9 @@ def control(*args, **kwargs):
         args[0] = [args[0]]
 
     if check_for_tracing_mode():
-        if all(isinstance(obj, AbstractQubit) for obj in [x.aval for x in args[0]]):
+        if all(isinstance(obj, bool) for obj in [x for x in args[0]]):
+            return ClControlEnvironment(*args, **kwargs)
+        elif all(isinstance(obj, AbstractQubit) for obj in [x.aval for x in args[0]]):
             return ControlEnvironment(*args, **kwargs)
         elif all(isinstance(obj, ShapedArray) for obj in [x.aval for x in args[0]]):
             return ClControlEnvironment(*args, **kwargs)

src/qrisp/jasp/control_flow/jrange_iterator.py

@@ -295,4 +295,11 @@ def make_tracer(x):
     def tracerizer():
         return jnp.array(x, dtype)
 
-    return jit(tracerizer)()
+    return jit(tracerizer)()
+
+def jlen(x):
+    if isinstance(x, list):
+        return len(x)
+    else:
+        return x.size
+    

src/qrisp/jasp/jasp_expression/centerclass.py

@@ -516,7 +516,13 @@ def qjit(self, *args, function_name = "jaspr_function"):
 
         from qrisp.jasp.catalyst_interface import jaspr_to_catalyst_qjit
         qjit_obj = jaspr_to_catalyst_qjit(flattened_jaspr, function_name = function_name)
-        return qjit_obj.compiled_function(*args)
+        res = qjit_obj.compiled_function(*args)
+        if not isinstance(res, (tuple,list)):
+            return res
+        elif len(res) == 1:
+            return res[0]
+        else:
+            return res
 
     @classmethod
     @lru_cache(maxsize = int(1E5))

tests/jax_tests/test_dynamic_mcx.py

@@ -38,4 +38,24 @@ def main(i, j):
                     assert k[1]
                 else:
                     assert not k[1]
+
+    # Test static list behavior
+    @terminal_sampling
+    def main(j):
+        qf = QuantumFloat(5)
+        h(qf)
+        qbl = QuantumBool()
+        qb_list = [qf[i] for i in range(5)]
+        mcx(qf.reg, qbl[0], method = "balauca", ctrl_state = j)
+        return qf, qbl
+
+    for j in range(2**5):
+        res_dict = main(j)
+        
+        for k in res_dict.keys():
+            if k[0] == j:
+                assert k[1]
+            else:
+                assert not k[1]
+