Fix import in benchmarks

This works on my machine as 'python benchmarks/pmap_benchmark.py'. It also
follows the code in examples.

This will need a copybara rule to change the import to 'from jax.benchmarks import benchmark'

benchmarks/pmap_benchmark.py

@@ -25,7 +25,7 @@ from jax import numpy as np
 from jax import pmap
 from jax.config import config
 
-import benchmark
+from benchmarks import benchmark
 
 import numpy as onp