Update Jax documentation

[rmoyard]

Context:

Correct some typos found by @ankit27kh

Related GitHub Issues:

Closes #1736

[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.

[codecov]

Codecov Report

Merging #1742 (7b95549) into master (c3c2e22) will not change coverage.
The diff coverage is n/a.

@@           Coverage Diff           @@
##           master    #1742   +/-   ##
=======================================
  Coverage   99.21%   99.21%           
=======================================
  Files         204      204           
  Lines       15467    15467           
=======================================
  Hits        15346    15346           
  Misses        121      121           

---

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 c3c2e22...7b95549. Read the comment docs.

[antalszava]

Looks good to me! 💯 Thanks for fixing this so quickly. 🙂

The warning likely comes from this line, where we work with the np.int64 type:
https://github.com/PennyLaneAI/pennylane/blob/master/pennylane/_qubit_device.py#L541