Qubit validation for Jasp IR

[ACE07-Sev]

Initial commit for tackling #232 .

Brief Overview

As noted in the slack discussion, this issue arises from missing checks on the jasp IR side which is trickier compared to the static IR side due to the dynamic nature of the SSA representation. Such checks include

• duplicate qubits
• mismatching number of qubits wrt the operation
• incompatible combination of qubit parameters for constructing qubit constellations

Such checks can only be done at simulation time, and this PR attempts to provide the checks at the most common intersection of the simulation functions.

Current commit allows for catching duplicate qubits, with the following MRE

from qrisp import *

@jaspify
def main():
    qv = QuantumBool()
    cx(qv, qv)
    return 0

main()

Exception: Duplicate qubit arguments in [Qubit(qb_60), Qubit(qb_60)] for operation cx

NOTE: Current commit does NOT check for valid qubit constellations given cases such as

from qrisp import *

@jaspify
def main():
    qv1 = QuantumFloat(3)
    qv2 = QuantumFloat(2)
    cx(qv1, qv2)
    return 0

main()

IndexError: list index out of range

---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
Cell In[3], [line 10](vscode-notebook-cell:?execution_count=3&line=10)
      7     cx(qv1, qv2)
      8     return 0
---> [10](vscode-notebook-cell:?execution_count=3&line=10) main()

File c:\Users\A.C.EA\Documents\GitHub\ucc_new_pass\.venv\Lib\site-packages\qrisp\jasp\evaluation_tools\jaspification.py:157, in jaspify.<locals>.return_function(*args)
    153     garbage_collection = "auto"
    154 jaspr = make_jaspr(tracing_function, garbage_collection=garbage_collection)(
    155     *args
    156 )
--> [157](file:///C:/Users/A.C.EA/Documents/GitHub/ucc_new_pass/.venv/Lib/site-packages/qrisp/jasp/evaluation_tools/jaspification.py:157) jaspr_res = simulate_jaspr(jaspr, *args, terminal_sampling=terminal_sampling)
    158 if isinstance(jaspr_res, tuple):
    159     jaspr_res = tree_unflatten(treedef_container[0], jaspr_res)

File c:\Users\A.C.EA\Documents\GitHub\ucc_new_pass\.venv\Lib\site-packages\qrisp\jasp\evaluation_tools\jaspification.py:352, in simulate_jaspr(jaspr, terminal_sampling, simulator, return_gate_counts, *args)
    349         return True
    351 with fast_append(3):
--> [352](file:///C:/Users/A.C.EA/Documents/GitHub/ucc_new_pass/.venv/Lib/site-packages/qrisp/jasp/evaluation_tools/jaspification.py:352)     res = eval_jaxpr(jaspr, eqn_evaluator=eqn_evaluator)(*(args))
    354 if return_gate_counts:
    355     return res[-1].gate_counts

File c:\Users\A.C.EA\Documents\GitHub\ucc_new_pass\.venv\Lib\site-packages\qrisp\jasp\interpreter_tools\abstract_interpreter.py:81, in eval_jaxpr.<locals>.jaxpr_evaluator(*args)
     78     raise Exception("Tried to evaluate jaxpr with insufficient arguments")
     80 context_dic = ContextDict({temp_var_list[i]: args[i] for i in range(len(args))})
---> [81](file:///C:/Users/A.C.EA/Documents/GitHub/ucc_new_pass/.venv/Lib/site-packages/qrisp/jasp/interpreter_tools/abstract_interpreter.py:81) eval_jaxpr_with_context_dic(jaxpr, context_dic, eqn_evaluator)
     83 if return_context_dic:
     84     outvals = [context_dic]

File c:\Users\A.C.EA\Documents\GitHub\ucc_new_pass\.venv\Lib\site-packages\qrisp\jasp\interpreter_tools\abstract_interpreter.py:141, in eval_jaxpr_with_context_dic(jaxpr, context_dic, eqn_evaluator)
    134 from qrisp.jasp import (
    135     evaluate_cond_eqn,
    136     evaluate_while_loop,
    137     evaluate_scan,
    138 )
    140 if eqn.primitive.name == "while":
--> [141](file:///C:/Users/A.C.EA/Documents/GitHub/ucc_new_pass/.venv/Lib/site-packages/qrisp/jasp/interpreter_tools/abstract_interpreter.py:141)     evaluate_while_loop(eqn, context_dic, eqn_evaluator)
    142 elif eqn.primitive.name == "cond":
    143     evaluate_cond_eqn(eqn, context_dic, eqn_evaluator)

File c:\Users\A.C.EA\Documents\GitHub\ucc_new_pass\.venv\Lib\site-packages\qrisp\jasp\interpreter_tools\interpreters\control_flow_interpretation.py:77, in evaluate_while_loop(while_loop_eqn, context_dic, eqn_evaluator, break_after_first_iter)
     73 carries = invalues[overall_constant_amount:]
     75 new_invalues = constants + carries
---> [77](file:///C:/Users/A.C.EA/Documents/GitHub/ucc_new_pass/.venv/Lib/site-packages/qrisp/jasp/interpreter_tools/interpreters/control_flow_interpretation.py:77) outvalues = eval_jaxpr(
     78     while_loop_eqn.params["body_jaxpr"], eqn_evaluator=eqn_evaluator
     79 )(*new_invalues)
     81 # Update the non-const invalues
     83 if len(while_loop_eqn.params["body_jaxpr"].jaxpr.outvars) == 1:

File c:\Users\A.C.EA\Documents\GitHub\ucc_new_pass\.venv\Lib\site-packages\qrisp\jasp\interpreter_tools\abstract_interpreter.py:81, in eval_jaxpr.<locals>.jaxpr_evaluator(*args)
     78     raise Exception("Tried to evaluate jaxpr with insufficient arguments")
     80 context_dic = ContextDict({temp_var_list[i]: args[i] for i in range(len(args))})
---> [81](file:///C:/Users/A.C.EA/Documents/GitHub/ucc_new_pass/.venv/Lib/site-packages/qrisp/jasp/interpreter_tools/abstract_interpreter.py:81) eval_jaxpr_with_context_dic(jaxpr, context_dic, eqn_evaluator)
     83 if return_context_dic:
     84     outvals = [context_dic]

File c:\Users\A.C.EA\Documents\GitHub\ucc_new_pass\.venv\Lib\site-packages\qrisp\jasp\interpreter_tools\abstract_interpreter.py:149, in eval_jaxpr_with_context_dic(jaxpr, context_dic, eqn_evaluator)
    145         evaluate_scan(eqn, context_dic, eqn_evaluator)
    147     continue
--> [149](file:///C:/Users/A.C.EA/Documents/GitHub/ucc_new_pass/.venv/Lib/site-packages/qrisp/jasp/interpreter_tools/abstract_interpreter.py:149) exec_eqn(eqn, context_dic)

File c:\Users\A.C.EA\Documents\GitHub\ucc_new_pass\.venv\Lib\site-packages\qrisp\jasp\interpreter_tools\abstract_interpreter.py:43, in exec_eqn(eqn, context_dic)
     41 def exec_eqn(eqn, context_dic):
     42     invalues = extract_invalues(eqn, context_dic)
---> [43](file:///C:/Users/A.C.EA/Documents/GitHub/ucc_new_pass/.venv/Lib/site-packages/qrisp/jasp/interpreter_tools/abstract_interpreter.py:43)     res = eqn.primitive.bind(*invalues, **eqn.params)
     44     insert_outvalues(eqn, context_dic, res)

File c:\Users\A.C.EA\Documents\GitHub\ucc_new_pass\.venv\Lib\site-packages\jax\_src\core.py:496, in Primitive.bind(self, *args, **params)
    494 def bind(self, *args, **params):
    495   args = args if self.skip_canonicalization else map(canonicalize_value, args)
--> [496](file:///C:/Users/A.C.EA/Documents/GitHub/ucc_new_pass/.venv/Lib/site-packages/jax/_src/core.py:496)   return self._true_bind(*args, **params)

File c:\Users\A.C.EA\Documents\GitHub\ucc_new_pass\.venv\Lib\site-packages\jax\_src\core.py:512, in Primitive._true_bind(self, *args, **params)
    510 trace_ctx.set_trace(eval_trace)
    511 try:
--> [512](file:///C:/Users/A.C.EA/Documents/GitHub/ucc_new_pass/.venv/Lib/site-packages/jax/_src/core.py:512)   return self.bind_with_trace(prev_trace, args, params)
    513 finally:
    514   trace_ctx.set_trace(prev_trace)

File c:\Users\A.C.EA\Documents\GitHub\ucc_new_pass\.venv\Lib\site-packages\jax\_src\core.py:517, in Primitive.bind_with_trace(self, trace, args, params)
    516 def bind_with_trace(self, trace, args, params):
--> [517](file:///C:/Users/A.C.EA/Documents/GitHub/ucc_new_pass/.venv/Lib/site-packages/jax/_src/core.py:517)   return trace.process_primitive(self, args, params)

File c:\Users\A.C.EA\Documents\GitHub\ucc_new_pass\.venv\Lib\site-packages\jax\_src\core.py:1017, in EvalTrace.process_primitive(self, primitive, args, params)
   1015 args = map(full_lower, args)
   1016 check_eval_args(args)
-> [1017](file:///C:/Users/A.C.EA/Documents/GitHub/ucc_new_pass/.venv/Lib/site-packages/jax/_src/core.py:1017) return primitive.impl(*args, **params)

File c:\Users\A.C.EA\Documents\GitHub\ucc_new_pass\.venv\Lib\site-packages\qrisp\jasp\primitives\abstract_quantum_register.py:147, in get_qubit_impl(qb_array, index)
    136 @get_qubit_p.def_impl
    137 def get_qubit_impl(qb_array, index):
    138     """Abstract evaluation of the primitive.
    139 
    140     This function does not need to be JAX traceable. It will be invoked with
   (...)    145       a ShapedArray for the result of the primitive.
    146     """
--> [147](file:///C:/Users/A.C.EA/Documents/GitHub/ucc_new_pass/.venv/Lib/site-packages/qrisp/jasp/primitives/abstract_quantum_register.py:147)     return qb_array[index]

IndexError: list index out of range

and instead, such cases either raise index error from `get_qubit_p` primitive or pass silently if `qv1` size is less than `qv2`. I will keep working to see if there is a way to add this as well.

[positr0nium]

The way you did it would slow the simulator down in general for a check that should only be necessary in Jasp mode. There is a simpler fix: Qrisp allows to trade such checks for speed using the fast_append context manager. There are several levels of fast append (0 - 3) which each deactivate different checks to accelerate the compilation. 3 is the fastest but also the most unstable.
The jasp simulator by default operates on level 3, which is why these checks don't happen. You can activate level 0 temporarily by modifying the code here:

with fast_append(0):
    self.buffer_qc.append(op, qubits)

fast_append can be imported from qrisp.circuit.

[ACE07-Sev]

Copy that, fixing it now.

[ACE07-Sev]

You can activate level 0 temporarily by modifying the code here:

Question. By temporarily, you mean this approach is not the way to go for solving this properly?

[positr0nium]

No, temporarily means that the ‘‘fast_append‘‘ mode is reset to its original mode after the context manager is left.