wip

graphcore-research · Nov 24, 2023 · 6aa5ed1 · 6aa5ed1
1 parent 18ae06b
commit 6aa5ed1
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 3 deletions.
diff --git a/experiments/mnist/mnist_classifier.py b/experiments/mnist/mnist_classifier.py
@@ -93,12 +93,17 @@ def update(i, opt_state, batch):
 
         # print(opt_state)
         for _ in range(num_batches):
-            batch = jax.tree_map(lambda v: jsa.scaled_array(v, scale=np.float32(1.0)), next(batches))
+            batch = next(batches)
+            batch = jax.tree_map(lambda v: jsa.scaled_array(v, scale=np.float32(1.0)), batch)
             # batch = jsa.scaled_array(next(batches), np.float32(0))
             opt_state = update(next(itercount), opt_state, batch)
         epoch_time = time.time() - start_time
 
         params = get_params(opt_state)
+
+        params = jax.tree_map(lambda v: jsa.core.asarray(v), params, is_leaf=jsa.core.is_scaled_leaf)
+        # print(params)
+
         train_acc = accuracy(params, (train_images, train_labels))
         test_acc = accuracy(params, (test_images, test_labels))
         print(f"Epoch {epoch} in {epoch_time:0.2f} sec")

diff --git a/jax_scaled_arithmetics/core/__init__.py b/jax_scaled_arithmetics/core/__init__.py
@@ -1,5 +1,13 @@
 # Copyright (c) 2023 Graphcore Ltd. All rights reserved.
-from .datatype import DTypeLike, ScaledArray, Shape, asarray, is_scaled_leaf, scaled_array  # noqa: F401
+from .datatype import (  # noqa: F401
+    DTypeLike,
+    ScaledArray,
+    Shape,
+    as_scaled_array,
+    asarray,
+    is_scaled_leaf,
+    scaled_array,
+)
 from .interpreters import (  # noqa: F401
     ScaledPrimitiveType,
     autoscale,

diff --git a/jax_scaled_arithmetics/core/datatype.py b/jax_scaled_arithmetics/core/datatype.py
@@ -103,6 +103,13 @@ def scaled_array(data: ArrayLike, scale: ArrayLike, dtype: DTypeLike = None, npa
     return ScaledArray(data, scale)
 
 
+def as_scaled_array(val: Any) -> ScaledArray:
+    """ """
+    if isinstance(val, ScaledArray):
+        return val
+    return ScaledArray(data=val, scale=np.array(1, dtype=val.dtype))
+
+
 def asarray(val: Any, dtype: DTypeLike = None) -> GenericArray:
     """Convert back to a common JAX/Numpy array.
 

diff --git a/jax_scaled_arithmetics/lax/scaled_ops.py b/jax_scaled_arithmetics/lax/scaled_ops.py
@@ -6,9 +6,10 @@
 import jax.numpy as jnp
 import numpy as np
 from jax import lax
+from jax._src.ad_util import add_any_p
 
 from jax_scaled_arithmetics import core
-from jax_scaled_arithmetics.core import DTypeLike, ScaledArray, Shape
+from jax_scaled_arithmetics.core import DTypeLike, ScaledArray, Shape, as_scaled_array, register_scaled_op
 
 from .base_scaling_primitives import scaled_set_scaling
 
@@ -87,6 +88,9 @@ def scaled_reshape(A: ScaledArray, new_sizes: Sequence[int], dimensions: Sequenc
 
 @core.register_scaled_lax_op
 def scaled_mul(A: ScaledArray, B: ScaledArray) -> ScaledArray:
+    # print("INPUTS:", A, B)
+    A = as_scaled_array(A)
+    B = as_scaled_array(B)
     return ScaledArray(A.data * B.data, A.scale * B.scale)
 
 
@@ -102,6 +106,9 @@ def scaled_add(A: ScaledArray, B: ScaledArray) -> ScaledArray:
     return ScaledArray(output_data, output_scale)
 
 
+register_scaled_op(add_any_p, scaled_add)
+
+
 @core.register_scaled_lax_op
 def scaled_sub(A: ScaledArray, B: ScaledArray) -> ScaledArray:
     check_scalar_scales(A, B)