PennyLaneAI / PennyLaneAI/catalyst

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

Open
#2,394 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug frontend
Dominant language
Python
Stars
234
Forks
84
Avg merge
2d 15h
Merged PRs (30d)
66

Description

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!

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Reproduce the failure using the issue's Python example with catalyst.jax_primitives.while_p and jax.make_jaxpr, then inspect the tracing path around DynshapePrimitive.bind and DynamicJaxprTrace.make_eqn. Done means the shown while_p.bind call no longer raises the unexpected out_tracers TypeError.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
compilers
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.