Improve API benchmarks.
Improve API benchmarks.

Add benchmarks for different dispatch arg arities.

Add more blocking before and after benchmark loops that don't otherwise block.

google/jax,google/jax,google/jax,google/jax