capture + dynamic shots + sample + applying a pass fails to create plxpr

[paul0403]

There is a mysterious failure when all of the following conditions are met:

• use capture
• apply a pass
• returns qml.sample() from a qnode
• use dynamic number of shots

When these conditions are met, JAX leaks an argument from the inner jaxpr of the pennylane transform primitive out into the returned outvar of the top level jaxpr, as the dynamically-shaped shot dimension for the samples.

See https://xanaduhq.slack.com/archives/C06CNPQLK2T/p1770325357111409

qml.capture.enable()
dev = qml.device("lightning.qubit", wires=2)

@qjit
def workflow(shots:int):
    @qml.transform(pass_name="cancel-inverses")
    @qml.qnode(dev, shots=shots)
    def aloha():
        qml.Hadamard(wires=0)
        return qml.sample()

    return aloha()

x = workflow(5)
print(x)

File "/home/paul.wang/catalyst_new/catalyst/frontend/catalyst/jit.py", line 735, in capture
    return trace_from_pennylane(
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/paul.wang/catalyst_new/catalyst/frontend/catalyst/from_plxpr/from_plxpr.py", line 536, in trace_from_pennylane
    jaxpr = from_plxpr(plxpr)(*plxpr.in_avals)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/paul.wang/catalyst_new/catalyst/frontend/catalyst/from_plxpr/from_plxpr.py", line 204, in wrapped_fn
    return jax.make_jaxpr(original_fn)(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/paul.wang/catalyst_new/catalyst/cat_env/lib/python3.12/site-packages/pennylane/capture/base_interpreter.py", line 389, in eval
    outval = self.read(var)
             ^^^^^^^^^^^^^^
  File "/home/paul.wang/catalyst_new/catalyst/cat_env/lib/python3.12/site-packages/pennylane/capture/base_interpreter.py", line 267, in read
    return var.val if isinstance(var, jax.extend.core.Literal) else self._env[var]
                                                                    ~~~~~~~~~^^^^^
KeyError: Var(id=134741294431552):int64[]

The cause is because the top-level jaxpr is trying to return a Var from an inner jaxpr

{ lambda ; a:i64[]. let
    b:i64[c,2] = transform[
      args_slice=(0, 0, None)
      consts_slice=(0, 1, None)
      inner_jaxpr={ lambda c:i64[]; . let
          d:i64[c,2] = qnode[
            device=<lightning.qubit device (wires=2) at 0x7f1a30552f00>
            execution_config=ExecutionConfig(grad_on_execution=False, use_device_gradient=None, use_device_jacobian_product=False, gradient_method='best', gradient_keyword_arguments={}, device_options={}, interface=<Interface.JAX: 'jax'>, derivative_order=1, mcm_config=MCMConfig(mcm_method=None, postselect_mode=None), convert_to_numpy=True, executor_backend=<class 'pennylane.concurrency.executors.native.multiproc.MPPoolExec'>)
            n_consts=1
            qfunc_jaxpr={ lambda e:i64[]; . let
                f:f64[] = convert_element_type[new_dtype=float64 weak_type=True] e
                g:f64[] = mul 0.1:f64[] f
                _:AbstractOperator() = RX[n_wires=1] g 1:i64[]
                _:AbstractOperator() = Hadamard[n_wires=1] 0:i64[]
                h:AbstractMeasurement(n_wires=0) = sample_wires[dtype=None] 
              in (h,) }
            qnode=<QNode: device='<lightning.qubit device (wires=2) at 0x7f1a30552f00>', interface='jax', diff_method='best', shots='Shots(total=JitTracer<~int64[]>)'>
            shots_len=1
          ] c c
        in (d,) }
      targs_slice=(1, None, None)
      tkwargs=()
      transform=<transform: cancel-inverses>
    ] a
  in (c, b) }

The outvar c is leaked from the inner jaxpr.

In jax/_src/interpreters/partial_eval.py/_add_implicit_outputs, the vars for the dynamic dimensions of a returned jaxpr value are added to the jaxpr outvars https://github.com/jax-ml/jax/blob/5712de44e97c455faed1fd45532e821ca66d025a/jax/_src/interpreters/partial_eval.py#L1843. Somehow for the transform primitive, the added Var is a leaked one from an inner jaxpr.

[paul0403]

Tried to patch away JAX's addition of the dynamic implicit dimension here PennyLaneAI/catalyst#2468

[paul0403]

Ok, after some digging, here's the root cause.

To make the plxpr, we call make_jaxpr2 in from_plxpr.py, which itself calls trace_to_jaxpr_dynamic2, a JAX internal function.

In trace_to_jaxpr_dynamic2, JAX calls _add_implicit_outputs, which takes in a jaxpr, and tries to append dynamic dimensions to the jaxpr outvars. It works like this:

1. Register all invars of the received jaxpr as "seen".
2. Look through all the explicit outvars of the received jaxpr. For each explicit outvar, do this:
   (They use a very nasty list comprehension; I asked gemini to unpack it)

impl_outvars = []
for x in expl_outvars:
    if type(x) is Var:
        seen.add(x)
        # Check if the variable has a dynamic shape
        if type(x.aval) is DShapedArray:
            for d in x.aval.shape:
                # If the dimension is itself a variable and new to us
                if type(d) is Var and d not in seen:
                    seen.add(d)
                    impl_outvars.append(d)

The important thing to realize is, they only add an implicit outvar if that dynamic dimension of the explicit outvar has not been "seen" before.

Why does this logic lead to the buggy behavior we are seeing? Well, look at the jaxpr received by this method, before appending the leaked var:

{ lambda ; a:i64[]. let
    b:i64[c,2] = transform[
      args_slice=(0, 0, None)
      consts_slice=(0, 1, None)
      inner_jaxpr={ lambda c:i64[]; . let
          d:i64[c,2] = qnode[
            device=<lightning.qubit device (wires=2) at 0x7319813faf30>
            execution_config=ExecutionConfig(grad_on_execution=False, use_device_gradient=None, use_device_jacobian_product=False, gradient_method='best', gradient_keyword_arguments={}, device_options={}, interface=<Interface.JAX: 'jax'>, derivative_order=1, mcm_config=MCMConfig(mcm_method=None, postselect_mode=None), convert_to_numpy=True, executor_backend=<class 'pennylane.concurrency.executors.native.multiproc.MPPoolExec'>)
            n_consts=0
            qfunc_jaxpr={ lambda ; . let
                _:AbstractOperator() = Hadamard[n_wires=1] 0:i64[]
                e:AbstractMeasurement(n_wires=0) = sample_wires[dtype=None] 
              in (e,) }
            qnode=<QNode: device='<lightning.qubit device (wires=2) at 0x7319813faf30>', interface='jax', diff_method='best', shots='Shots(total=JitTracer<~int64[]>)'>
            shots_len=1
          ] c
        in (d,) }
      targs_slice=(1, None, None)
      tkwargs=()
      transform=<transform: cancel-inverses>
    ] a
  in (b,) }  # no leak yet

The explicit outvar b has shape [c,2]. Because c is not an invar to the jaxpr (only a is), it is not considered "seen", so it got appended to the jaxpr outvars by this _add_implicit_outputs method.

By constrast, let's look at what happens when we don't apply a pass, and hence not binding the PL transform primitive. In this case, the _add_implicit_outputs method receives this jaxpr:

{ lambda ; a:i64[]. let
    b:i64[a,2] = qnode[
      device=<lightning.qubit device (wires=2) at 0x7ca08af6aea0>
      execution_config=ExecutionConfig(grad_on_execution=False, use_device_gradient=None, use_device_jacobian_product=False, gradient_method='best', gradient_keyword_arguments={}, device_options={}, interface=<Interface.JAX: 'jax'>, derivative_order=1, mcm_config=MCMConfig(mcm_method=None, postselect_mode=None), convert_to_numpy=True, executor_backend=<class 'pennylane.concurrency.executors.native.multiproc.MPPoolExec'>)
      n_consts=0
      qfunc_jaxpr={ lambda ; . let
          _:AbstractOperator() = Hadamard[n_wires=1] 0:i64[]
          c:AbstractMeasurement(n_wires=0) = sample_wires[dtype=None] 
        in (c,) }
      qnode=<QNode: device='<lightning.qubit device (wires=2) at 0x7ca08af6aea0>', interface='jax', diff_method='best', shots='Shots(total=JitTracer<~int64[]>)'>
      shots_len=1
    ] a
  in (b,) }

The explicit outvar b still is dynamically shaped, but its dynamic dimension a is now an invar to the top-level jaxpr, so it is considered "seen", and never got appended to the jaxpr outvars! Therefore, the leak does not happen in this case.

Now, I think this is because when determining the outshape of samples, the shots dimension is registered only once at the creation of the abatrct measurement, so this dimension when shots is dynamic is forever tied to the tracer for shots in the jaxpr that has the sample primitive. When that jaxpr becomes an inner jaxpr of some other primitive (transform primitive in this case), the sample shape still thinks shots is the inner jaxpr's shots argument tracer/var, and never promoted it to the shots argument of the outer jaxpr.