wip

Author: balancap

experiments/mnist/mnist_classifier_from_scratch.py

@@ -17,7 +17,6 @@
 The primary aim here is simplicity and minimal dependencies.
 """
 import time
-from functools import partial
 
 import datasets
 import jax
@@ -29,6 +28,8 @@
 
 import jax_scaled_arithmetics as jsa
 
+# from functools import partial
+
 
 def print_mean_std(name, v):
     data, scale = jsa.lax.get_data_scale(v)
@@ -60,7 +61,8 @@ def predict(params, inputs):
     # jsa.ops.debug_callback(partial(print_mean_std, "Logits"), logits)
     # (logits,) = jsa.ops.debug_callback_grad(partial(print_mean_std, "LogitsGrad"), logits)
 
-    logits = jsa.ops.dynamic_rescale_l1_grad(logits)
+    logits = jsa.ops.dynamic_rescale_l2_grad(logits)
+    # logits = logits.astype(np.float32)
     # (logits,) = jsa.ops.debug_callback_grad(partial(print_mean_std, "LogitsGrad"), logits)
 
     return logits - logsumexp(logits, axis=1, keepdims=True)
@@ -110,7 +112,7 @@ def data_stream():
     @jsa.autoscale
     def update(params, batch):
         grads = grad(loss)(params, batch)
-        # return [(w - step_size * dw, b - step_size * db) for (w, b), (dw, db) in zip(params, grads)]
+        return [(w - step_size * dw, b - step_size * db) for (w, b), (dw, db) in zip(params, grads)]
         return [
             (jsa.ops.dynamic_rescale_l1(w - step_size * dw), jsa.ops.dynamic_rescale_l1(b - step_size * db))
             for (w, b), (dw, db) in zip(params, grads)