dynamic_one_shot uses tapes with shot-vectors and jitting takes advantage of it

[vincentmr]

Before submitting

Please complete the following checklist when submitting a PR:

• All new features must include a unit test.
  If you've fixed a bug or added code that should be tested, add a test to the
  test directory!

• All new functions and code must be clearly commented and documented.
  If you do make documentation changes, make sure that the docs build and
  render correctly by running make docs.

• Ensure that the test suite passes, by running make test.

• Add a new entry to the doc/releases/changelog-dev.md file, summarizing the
  change, and including a link back to the PR.

• The PennyLane source code conforms to
  PEP8 standards.
  We check all of our code against Pylint.
  To lint modified files, simply pip install pylint, and then
  run pylint pennylane/path/to/file.py.

When all the above are checked, delete everything above the dashed
line and fill in the pull request template.

---

Context:
dynamic_one_shot creates n-shots tapes which is wasteful.

Description of the Change:
Create a single tape with a shot-vector which indicates to the device how many times to repeat the tape execution.

Benefits:
For a tape like

dev = qml.device("default.qubit", shots=2000, seed=jax.random.PRNGKey(123))

@qml.qnode(dev, diff_method=None)
def func(x, y):
    qml.RX(x, wires=0)
    m0 = qml.measure(0, reset=False, postselect=1)
    qml.cond(m0, qml.RY)(y, wires=1)
    return qml.expval(qml.PauliZ(0))


params = np.pi / 4 * np.ones(2)

The execution times are as follows (Latitude laptop):

• 12.7 s : vanilla Python
• 8.2 s : jax.vmap
• 11.1 s : jax.jit + jax.vmap + compilation
• 6.89 ms : jax.jit + jax.vmap

Possible Drawbacks:

Related GitHub Issues:
[sc-62097]

[dime10]

This looks good to me 👍

[albi3ro]

Changelog?

[albi3ro]

Might be worth investigating if we can somehow use use multi-processing with this in the future.

Thanks for doing all this investigation!

[vincentmr]

Might be worth investigating if we can somehow use use multi-processing with this in the future.

That's actually the original idea behind having multiple tapes instead of shots. On the other hand, I think multi-processing might be better turned-on at the QNode level, or maybe with a decorator. I think any device with a batch of tapes could be executing them in parallel inside sub-processes.

[dime10]

Might be worth investigating if we can somehow use use multi-processing with this in the future.

That's actually the original idea behind having multiple tapes instead of shots. On the other hand, I think multi-processing might be better turned-on at the QNode level, or maybe with a decorator. I think any device with a batch of tapes could be executing them in parallel inside sub-processes.

Why couldn't you parallelize with shots instead of tapes?

[vincentmr]

Why couldn't you parallelize with shots instead of tapes?

The current multiprocessing logic live here. It is not a fundamental limitation, it's just our UI is not set up for it.