basic interpeter for scaled ops and scaled arrays (#8)

* basic interpeter for scaled ops and scaled arrays

* precommit

* mypy fixes

* use ScaledArray.aval

* move registry into core

* Fixing module imports

* imports in top level __init__.py

* linting fixes

* return value if in singleton list

* return value if in singleton list

* comment on reason for multiple vars per scaledarray

---------

Co-authored-by: Paul Balanca <paulb@graphcore.ai>

jax_scaled_arithmetics/__init__.py

@@ -1,2 +1,4 @@
 # Copyright (c) 2023 Graphcore Ltd. All rights reserved.
+from . import lax
 from ._version import __version__
+from .core import ScaledArray, autoscale # noqa: F401

jax_scaled_arithmetics/core/__init__.py

@@ -1,2 +1,3 @@
 # Copyright (c) 2023 Graphcore Ltd. All rights reserved.
 from .datatype import ScaledArray # noqa: F401
+from .interpreters import autoscale, register_scaled_op # noqa: F401

jax_scaled_arithmetics/core/datatype.py

@@ -78,3 +78,7 @@ def to_array(self, dtype: DTypeLike = None) -> GenericArray:
  def __array__(self, dtype: DTypeLike = None) -> NDArray[Any]:
  """Numpy array interface support."""
  return np.asarray(self.to_array(dtype))
+
+ @property
+ def aval(self):
+ return self.data * self.scale

jax_scaled_arithmetics/core/interpreters.py

@@ -0,0 +1,59 @@
+# Copyright (c) 2023 Graphcore Ltd. All rights reserved.
+
+from functools import wraps
+from typing import Dict
+
+import jax
+from jax import core
+from jax._src.util import safe_map
+
+from ..core import ScaledArray
+
+_scaled_ops_registry = {}
+
+
+def register_scaled_op(lax_func, scaled_func):
+ _scaled_ops_registry[lax_func] = scaled_func
+
+
+def autoscale(fun):
+ @wraps(fun)
+ def wrapped(*args, **kwargs):
+ aval_args = safe_map(lambda x: x.aval, args)
+ # get jaxpr of unscaled graph
+ closed_jaxpr = jax.make_jaxpr(fun)(*aval_args, **kwargs)
+ # convert to scaled graph
+ out = autoscale_jaxpr(closed_jaxpr.jaxpr, closed_jaxpr.literals, *args)
+ return out
+
+ return wrapped
+
+
+def autoscale_jaxpr(jaxpr, consts, *args):
+ env: Dict[core.Var, ScaledArray] = {}
+
+ def read(var):
+ if type(var) is core.Literal:
+ return var.val
+ return env[var]
+
+ def write(var, val):
+ env[var] = val
+
+ safe_map(write, jaxpr.invars, args)
+ safe_map(write, jaxpr.constvars, consts)
+
+ for eqn in jaxpr.eqns:
+ invals = safe_map(read, eqn.invars)
+ if eqn.primitive not in _scaled_ops_registry:
+ raise NotImplementedError(f"{eqn.primitive} does not have an implementation for ScaledArray inputs yet")
+ outvals = _scaled_ops_registry[eqn.primitive](*invals)
+ if not eqn.primitive.multiple_results:
+ outvals = [outvals]
+ safe_map(write, eqn.outvars, outvals)
+
+ outvals = safe_map(read, jaxpr.outvars)
+ if len(outvals) == 1:
+ return outvals[0]
+ else:
+ return outvals

jax_scaled_arithmetics/lax/__init__.py

@@ -1 +1,2 @@
 # Copyright (c) 2023 Graphcore Ltd. All rights reserved.
+from .scaled_ops import * # noqa: F401, F403

jax_scaled_arithmetics/lax/scaled_ops.py

@@ -0,0 +1,15 @@
+# Copyright (c) 2023 Graphcore Ltd. All rights reserved.
+
+from jax import lax
+
+from jax_scaled_arithmetics import core
+from jax_scaled_arithmetics.core import ScaledArray
+
+
+def scaled_mul_p(A: ScaledArray, B: ScaledArray) -> ScaledArray:
+ return ScaledArray(A.data * B.data, A.scale * B.scale)
+
+
+core.register_scaled_op(lax.mul_p, scaled_mul_p)
+
+__all__ = ["scaled_mul_p"]

tests/interpreters/test_interpreter.py

@@ -0,0 +1,54 @@
+# Copyright (c) 2023 Graphcore Ltd. All rights reserved.
+
+import chex
+import jax
+import jax.numpy as jnp
+
+from jax_scaled_arithmetics.core import ScaledArray, autoscale
+
+
+class AutoScaleInterpreterTests(chex.TestCase):
+ def test__identity(self):
+ def func(x):
+ return x
+
+ asfunc = autoscale(func)
+
+ scale = jnp.array(1.0)
+ inputs = jnp.array([1.0, 2.0])
+ expected = jnp.array([1.0, 2.0])
+
+ scaled_inputs = ScaledArray(inputs, scale)
+ scaled_outputs = asfunc(scaled_inputs)
+
+ assert jnp.allclose(scaled_outputs.aval, expected)
+
+ jaxpr = jax.make_jaxpr(asfunc)(scaled_inputs).jaxpr
+
+ # Vars need to be primitive data types (e.g., f32) -> 2 Vars per ScaledArray
+
+ assert jaxpr.invars[0].aval.shape == inputs.shape
+ assert jaxpr.invars[1].aval.shape == ()
+
+ assert jaxpr.outvars[0].aval.shape == expected.shape
+ assert jaxpr.outvars[1].aval.shape == ()
+
+ def test__mul(self):
+ def func(x, y):
+ return x * y
+
+ asfunc = autoscale(func)
+
+ x_in = jnp.array([-2.0, 2.0])
+ x_scale = jnp.array(0.5)
+ x = ScaledArray(x_in, x_scale)
+
+ y_in = jnp.array([1.5, 1.5])
+ y_scale = jnp.array(2.0)
+ y = ScaledArray(y_in, y_scale)
+
+ expected = jnp.array([-3.0, 3.0])
+
+ out = asfunc(x, y)
+
+ assert jnp.allclose(out.aval, expected)