mcm_method='one-shot' is not supported with a dynamic number of shots

[dime10]

Possibly related to a similar issue with dynamic shapes in #1949, the dynamic one-shot transform does not work when the number of shots is a tracer value. The underlying issue is the vmap transform not being able to handle a dynamic batch size.

import pennylane as qml

@qml.qjit
def main(shots: int):

    @qml.qnode(qml.device("lightning.qubit", wires=1), mcm_method="one-shot")
    def f():
        qml.H(0)
        return qml.sample()

    return qml.set_shots(shots)(f)()

main(10)

  File "catalyst/frontend/catalyst/api_extensions/function_maps.py", line 368, in _get_batch_size
    if not batch_size:

jax.errors.TracerBoolConversionError: Attempted boolean conversion of traced array with shape int64[].