wip

Author: balancap

experiments/mnist/mnist_classifier_from_scratch.py

@@ -31,10 +31,11 @@
 
 
 def print_mean_std(name, v):
+    _, scale = jsa.lax.get_data_scale(v)
     # Always use np.float32, to avoid floating errors in descaling + stats.
     v = jsa.asarray(v, dtype=np.float32)
     m, s = np.mean(v), np.std(v)
-    print(name, m, s)
+    print(f"{name}: MEAN({m:.4f}) / STD({s:.4f}) / SCALE({scale:.4f})")
 
 
 def init_random_params(scale, layer_sizes, rng=npr.RandomState(0)):
@@ -44,8 +45,8 @@ def init_random_params(scale, layer_sizes, rng=npr.RandomState(0)):
 def predict(params, inputs):
     activations = inputs
     for w, b in params[:-1]:
-        jsa.ops.debug_callback(partial(print_mean_std, "W:"), w)
-        (w,) = jsa.ops.debug_callback_grad(partial(print_mean_std, "WG:"), w)
+        jsa.ops.debug_callback(partial(print_mean_std, "W"), w)
+        (w,) = jsa.ops.debug_callback_grad(partial(print_mean_std, "WGrad"), w)
 
         # Matmul + relu
         outputs = jnp.dot(activations, w) + b
@@ -54,8 +55,8 @@ def predict(params, inputs):
     final_w, final_b = params[-1]
     logits = jnp.dot(activations, final_w) + final_b
 
-    jsa.ops.debug_callback(partial(print_mean_std, "LOGITS:"), logits)
-    (logits,) = jsa.ops.debug_callback_grad(partial(print_mean_std, "LOGITSG:"), logits)
+    jsa.ops.debug_callback(partial(print_mean_std, "Logits"), logits)
+    (logits,) = jsa.ops.debug_callback_grad(partial(print_mean_std, "LogitsGrad"), logits)
 
     return logits - logsumexp(logits, axis=1, keepdims=True)