PennyLaneAI / PennyLaneAI/catalyst
Calling ``while_p.bind`` results in an error
Nobody has claimed this yet.
- 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
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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