Speed benchmarking/profiling
It might also be useful to have a speed benchmark, so that it's clear whether a change slows or speeds up the code. For example, replacing `uu/jnp.sqrt(jnp.sum(jnp.square(uu)))` with `uu/jnp.linalg.norm(uu)` seems innocent, but may slow down code (https://stackoverflow.com/questions/64948677/why-is-np-linalg-norm-axis-1-slower-than-writing-out-the-formula-for-vecto), and this sort of thing would be good to catch automatically. Probably just running the sampler on a multivariate gaussian is fine, since we're interested here in computational efficiency, not statistical.
The usual way this works, for pure deterministic functions (which this is, thanks to Jax) is that you run the program some large number of times, and extract statistics about mean speed and variance. There are various Python libraries that do this for you.
1 条评论