Fix JAX qml.sample example in the docs

[antalszava]

As per https://discuss.pennylane.ai/t/pennylane-0-20-qml-sample-not-jittable-anymore/1664

[github-actions]

Hello. You may have forgotten to update the changelog!
Please edit doc/releases/changelog-dev.md with:

• A one-to-two sentence description of the change. You may include a small working example for new features.
• A link back to this PR.
• Your name (or GitHub username) in the contributors section.

[josh146]

Thanks Antal :)

[codecov]

Codecov Report

Merging #2196 (c9cb7cc) into master (4b6155a) will not change coverage.
The diff coverage is n/a.

@@           Coverage Diff           @@
##           master    #2196   +/-   ##
=======================================
  Coverage   99.21%   99.21%           
=======================================
  Files         231      231           
  Lines       17901    17901           
=======================================
  Hits        17760    17760           
  Misses        141      141           

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 4b6155a...c9cb7cc. Read the comment docs.