wip

Author: balancap

experiments/mnist/mnist_classifier_from_scratch.py

@@ -45,13 +45,18 @@ 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)
 
         # Matmul + relu
         outputs = jnp.dot(activations, w) + b
         activations = jnp.maximum(outputs, 0)
 
     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)
+
     return logits - logsumexp(logits, axis=1, keepdims=True)