We read every piece of feedback, and take your input very seriously.
To see all available qualifiers, see our documentation.
There was an error while loading. Please reload this page.
1 parent 15755c9 commit 705f963Copy full SHA for 705f963
pennylane/capture/dynamic_shapes.py
@@ -232,16 +232,7 @@ def custom_staging_rule(
232
jax.core.no_effects,
233
source_info,
234
)
235
- tracing_eqn = TracingEqn(
236
- list(tracers),
237
- returned_vars,
238
- primitive,
239
- params,
240
- eqn.effects,
241
- source_info,
242
- eqn.ctx,
243
- )
244
- jaxpr_trace.frame.add_eqn(tracing_eqn)
+ jaxpr_trace.frame.add_eqn(eqn)
245
return out_tracers
246
247
pe.custom_staging_rules[primitive] = custom_staging_rule
0 commit comments