Fix MNIST training example with dynamic L2 rescale logits. (#76)

Training set accuracy 0.96670
Test set accuracy 0.93730

Note: test set accuracy can change slightly if using an `eps` in dynamic rescaling.

Author: balancap

experiments/mnist/mnist_classifier_from_scratch.py

@@ -44,6 +44,8 @@ def predict(params, inputs):
 
     final_w, final_b = params[-1]
     logits = jnp.dot(activations, final_w) + final_b
+    # Dynamic rescaling of the gradient, as logits gradient not properly scaled.
+    logits = jsa.ops.dynamic_rescale_l2_grad(logits)
     return logits - logsumexp(logits, axis=1, keepdims=True)