        'accuracy': accuracy,
    }
 +  if sharded:
-+    metrics = pmean(metrics)
++    metrics = lax.pmean(metrics, 'batch')
    return metrics
  
  

 def train_step(optimizer, batch):
   """Train for a single step."""
   def loss_fn(model):
@@ -121,18 +139,23 @@ def train_step(optimizer, batch):
   grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
   (_, logits), grad = grad_fn(optimizer.target)
   optimizer = optimizer.apply_gradient(grad)
-  metrics = compute_metrics(logits, batch['label'])
+  metrics = compute_metrics(logits, batch['label'], sharded=True)

diff --git a/examples/mnist/train.py b/examples/mnist/train.py
index 51d2fde..08951b7 100644
--- a/examples/mnist/train.py
+++ b/examples/mnist/train.py
@@ -23,10 +23,14 @@ from absl import app
 from absl import flags
 from absl import logging
 
+import functools
+
+from flax import jax_utils
 from flax import nn
 from flax import optim
 
 import jax
+from jax import lax
 from jax import random
 
 import jax.numpy as jnp
@@ -89,6 +93,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 = jax_utils.replicate(optimizer)
   return optimizer
 
 
@@ -101,17 +106,28 @@ def cross_entropy_loss(logits, labels):
   return -jnp.mean(jnp.sum(onehot(labels) * logits, axis=-1))
 
 
