fixed an issue that prevent the proper conversion of literals to i32 within custom_control

Author: positr0nium

src/qrisp/environments/custom_control_environment.py

@@ -19,6 +19,7 @@
 import inspect
 
 import jax
+import jax.numpy as jnp
 
 from qrisp.environments.quantum_environments import QuantumEnvironment
 from qrisp.environments.gate_wrap_environment import GateWrapEnvironment
@@ -206,6 +207,19 @@ def adaptive_control_function(*args, **kwargs):
                     res = func(*args, ctrl = control_qb, **kwargs)
 
         else:
+            
+            args = list(args)
+            if func.__name__ == "extract_boolean_digit":
+                print(args)
+            for i in range(len(args)):
+                if isinstance(args[i], bool):
+                    args[i] = jnp.array(args[i], dtype = jnp.bool)
+                elif isinstance(args[i], int):
+                    args[i] = jnp.array(args[i], dtype = jnp.int32)
+                elif isinstance(args[i], float):
+                    args[i] = jnp.array(args[i], dtype = jnp.float32)
+                elif isinstance(args[i], complex):
+                    args[i] = jnp.array(args[i], dtype = jnp.complex)
 
             # Call the (qached) function
             res = func(*args, **kwargs)

src/qrisp/jasp/tracing_logic/qaching.py

@@ -231,18 +231,6 @@ def ammended_function(abs_qc, *args, **kwargs):
             abs_qs.register_qv(qv, None)
             flattened_qvs.extend(list(flatten_qv(qv)[0]))
 
-        # Make sure literals are 32 bit
-        args = list(args)
-        for i in range(len(args)):
-            if isinstance(args[i], bool):
-                args[i] = jnp.array(args[i], dtype = jnp.bool)
-            elif isinstance(args[i], int):
-                args[i] = jnp.array(args[i], dtype = jnp.int32)
-            elif isinstance(args[i], float):
-                args[i] = jnp.array(args[i], dtype = jnp.float32)
-            elif isinstance(args[i], complex):
-                args[i] = jnp.array(args[i], dtype = jnp.complex)
-        
         # Execute the function
         res = func(*args, **kwargs)
         new_abs_qc = abs_qs.abs_qc
@@ -283,6 +271,18 @@ def return_function(*args, **kwargs):
         # Get the AbstractQuantumCircuit for tracing
         abs_qs = TracingQuantumSession.get_instance()
         abs_qs.start_tracing(abs_qs.abs_qc)
+
+        # Make sure literals are 32 bit
+        args = list(args)
+        for i in range(len(args)):
+            if isinstance(args[i], bool):
+                args[i] = jnp.array(args[i], dtype = jnp.bool)
+            elif isinstance(args[i], int):
+                args[i] = jnp.array(args[i], dtype = jnp.int32)
+            elif isinstance(args[i], float):
+                args[i] = jnp.array(args[i], dtype = jnp.float32)
+            elif isinstance(args[i], complex):
+                args[i] = jnp.array(args[i], dtype = jnp.complex)
 
         # Excecute the function
         abs_qc_new, res = ammended_function(abs_qs.abs_qc, *args, **kwargs)