Revert import path changes to examples/ and benchmarks/

PiperOrigin-RevId: 352911869

Author: hawkinsp

benchmarks/pmap_benchmark.py

@@ -26,7 +26,7 @@
 from jax.config import config
 from jax._src.util import prod
 
-from . import benchmark
+from benchmarks import benchmark
 
 import numpy as np
 

examples/mnist_classifier.py

@@ -30,7 +30,7 @@
 from jax.experimental import optimizers
 from jax.experimental import stax
 from jax.experimental.stax import Dense, Relu, LogSoftmax
-from . import datasets
+from examples import datasets
 
 
 def loss(params, batch):

examples/mnist_classifier_fromscratch.py

@@ -22,10 +22,10 @@
 
 import numpy.random as npr
 
-from jax.api import jit, grad
+from jax import jit, grad
 from jax.scipy.special import logsumexp
 import jax.numpy as jnp
-from . import datasets
+from examples import datasets
 
 
 def init_random_params(scale, layer_sizes, rng=npr.RandomState(0)):

examples/mnist_vae.py

@@ -30,7 +30,7 @@
 from jax.experimental import optimizers
 from jax.experimental import stax
 from jax.experimental.stax import Dense, FanOut, Relu, Softplus
-from . import datasets
+from examples import datasets
 
 
 def gaussian_kl(mu, sigmasq):