 +  return jax.device_get(jax.tree_map(lambda x: x[0], xs))
 +
 +
-+def pmean(tree, axis_name='batch'):
-+  return jax.tree_map(
-+      lambda x: lax.psum(x, axis_name) / jax.device_count(), tree)
-+
-+
 +def compute_metrics(logits, labels, sharded):
    loss = cross_entropy_loss(logits, labels)
    accuracy = jnp.mean(jnp.argmax(logits, -1) == labels)

diff --git a/examples/mnist/train.py b/examples/mnist/train.py
index 51d2fde..082faaa 100644
--- a/examples/mnist/train.py
+++ b/examples/mnist/train.py
@@ -26,7 +26,8 @@ from absl import logging
 from flax import nn
 from flax import optim
 
-import jax
+import 
+from jax import lax
 from jax import random
 
 import jax.numpy as jnp
@@ -47,8 +48,8 @@ flags.DEFINE_float(
     help=('The decay rate used for the momentum optimizer.'))
 
 flags.DEFINE_integer(
-    'batch_size', default=128,
-    help=('Batch size for training.'))
+    'batch_size', default=128 * jax.device_count(),
+    help=('Batch size for distributed single-host training.'))
 
 flags.DEFINE_integer(
     'num_epochs', default=10,
@@ -89,6 +90,7 @@ def create_model(key):
 def create_optimizer(model, learning_rate, beta):
   optimizer_def = optim.Momentum(learning_rate=learning_rate, beta=beta)
   optimizer = optimizer_def.create(model)
+  optimizer = optimizer.replicate()
   return optimizer
 
 
@@ -101,17 +103,33 @@ def cross_entropy_loss(logits, labels):
   return -jnp.mean(jnp.sum(onehot(labels) * logits, axis=-1))
 
 

 def train_step(optimizer, batch):
   """Train for a single step."""
   def loss_fn(model):
@@ -120,19 +136,25 @@ def train_step(optimizer, batch):
     return loss, logits
   grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
   (_, logits), grad = grad_fn(optimizer.target)
+  grad = lax.pmean(grad, 'batch')
   optimizer = optimizer.apply_gradient(grad)
-  metrics = compute_metrics(logits, batch['label'])
+  metrics = compute_metrics(logits, batch['label'], sharded=True)
