fixed cross-entropy losses in mnist examples (fixes #1023)

examples/differentially_private_sgd.py

@@ -127,7 +127,7 @@ def loss(params, batch):
   inputs, targets = batch
   logits = predict(params, inputs)
   logits = stax.logsoftmax(logits)  # log normalize
-  return -np.mean(np.sum(logits * targets, 1))  # cross entropy loss
+  return -np.mean(np.sum(logits * targets, axis=1))  # cross entropy loss
 
 
 def accuracy(params, batch):

examples/mnist_classifier.py

@@ -40,7 +40,7 @@ from examples import datasets
 def loss(params, batch):
   inputs, targets = batch
   preds = predict(params, inputs)
-  return -np.mean(preds * targets)
+  return -np.mean(np.sum(preds * targets, axis=1))
 
 def accuracy(params, batch):
   inputs, targets = batch

examples/mnist_classifier_fromscratch.py

@@ -49,7 +49,7 @@ def predict(params, inputs):
 def loss(params, batch):
   inputs, targets = batch
   preds = predict(params, inputs)
-  return -np.mean(preds * targets)
+  return -np.mean(np.sum(preds * targets, axis=1))
 
 def accuracy(params, batch):
   inputs, targets = batch

examples/spmd_mnist_classifier_fromscratch.py

@@ -57,7 +57,7 @@ def predict(params, inputs):
 def loss(params, batch):
   inputs, targets = batch
   preds = predict(params, inputs)
-  return -np.mean(preds * targets)
+  return -np.mean(np.sum(preds * targets, axis=1))
 
 @jit
 def accuracy(params, batch):