Tests: benchmarks

[nstarman]

Followup to #67.

Requires:

1. Codspeed to be set up
2. the tag run-benchmarks to be added.

[nstarman]

@patrick-kidger I think this PR may be ready.
If you could hook quax up to codspeed then we can maybe try it out. Alternatively, if you don't want codspeed I can try refactoring this for vanilla pytest-benchmark.

[johannahaffner]

I had a longer discussion with @jpbrodrick on this topic, codspeed and JAX really do not seem to mesh well.
(Discussion started here.)

What I have settled on is using just pytest-benchmark, and running these locally (not in CI).