Calling ``while_p.bind`` results in an error

[positr0nium]

Hi Catalyst team and congratulations to your new release!

When migrating to the latest version, I came across an issue that reproduced by the following code:

from jax import make_jaxpr
from catalyst.jax_primitives import AbstractQreg, while_p

def body(val):
    return val

body_jaxpr = make_jaxpr(body)(AbstractQreg())

def cond(val):
    return True

cond_jaxpr = make_jaxpr(body)(AbstractQreg())

def main(qreg):

    while_res = while_p.bind(
        qreg,
        cond_jaxpr=cond_jaxpr,
        body_jaxpr=body_jaxpr,
        cond_nconsts=0,
        body_nconsts=0,
        nimplicit=0,
        preserve_dimensions=True,
        num_implicit_inputs=0
    )
    
res = make_jaxpr(main)(AbstractQreg())

This results in

TypeError                                 Traceback (most recent call last)
File ~/.config/spyder-py3/temp.py:2802
   2800     )
   2801     
-> 2802 res = make_jaxpr(main)(AbstractQreg())

    [... skipping hidden 14 frame]

File ~/.config/spyder-py3/temp.py:2791, in main(qreg)
   2789 def main(qreg):
   2790 
-> 2791     while_res = while_p.bind(
   2792         qreg,
   2793         cond_jaxpr=cond_jaxpr,

File ~/miniforge3/envs/qrisp/lib/python3.11/site-packages/catalyst/jax_extras/tracing.py:949, in DynshapePrimitive.bind(self, *args, **params)
    945 # invars = map(lambda t: t.val, tracers)
    946 # outvars = map(lambda t: t.val, out_tracers)
    947 # eqn = new_jaxpr_eqn(invars, outvars, self, params, [], source_info)
    948 out_avals = [t.aval for t in out_tracers]
--> 949 eqn, out_tracers = trace.make_eqn(
    950     tracers, out_avals, self, params, [], source_info, out_tracers=out_tracers
    951 )
    952 trace.frame.add_eqn(eqn)
    953 return out_tracers if self.multiple_results else out_tracers.pop()

TypeError: DynamicJaxprTrace.make_eqn() got an unexpected keyword argument 'out_tracers'

If you know any way for me to circumvent this, I could already upgrade Qrisp to the latest Catalyst version. Thanks for your work!

[dime10]

Thanks for reporting this @positr0nium! This is due to some patching we had to do after upgrading to a recent JAX version, where they have abandoned dynamic shape support entirely. We're working on more solid options, but if you want to get this working right now you probably just need to ensure our patch is active when tracing (it would normally be active when using qjit, but since that's not the case here):

from catalyst.jax_extras.patches import patched_make_eqn
from catalyst.utils.patching import Patcher
from jax.interpreters.partial_eval import DynamicJaxprTrace

with Patcher(
    (DynamicJaxprTrace, "make_eqn", patched_make_eqn),
):
    make_jaxpr(main)(AbstractQreg())