Add SR on arrays, add visualization on notebook

Author: awf

docs/source/05-stochastic-rounding.ipynb

@@ -1,5 +1,6 @@
 # Copyright (c) 2024 Graphcore Ltd. All rights reserved.
 
+from typing import Optional
 from types import ModuleType
 from .types import FormatInfo, RoundMode
 import numpy as np
@@ -15,6 +16,8 @@ def round_ndarray(
     v: np.ndarray,
     rnd: RoundMode = RoundMode.TiesToEven,
     sat: bool = False,
+    srbits: Optional[np.ndarray] = None,
+    srnumbits: int = 0,
     np: ModuleType = np,
 ) -> np.ndarray:
     """
@@ -70,18 +73,24 @@ def round_ndarray(
     else:
         code_is_odd = (isignificand != 0) & _isodd(expval + bias)
 
+    if rnd == RoundMode.TowardZero:
+        should_round_away = np.zeros_like(delta, dtype=bool)
     if rnd == RoundMode.TowardPositive:
-        round_up = ~is_negative & (delta > 0)
-    elif rnd == RoundMode.TowardNegative:
-        round_up = is_negative & (delta > 0)
-    elif rnd == RoundMode.TiesToAway:
-        round_up = delta >= 0.5
-    elif rnd == RoundMode.TiesToEven:
-        round_up = (delta > 0.5) | ((delta == 0.5) & code_is_odd)
-    else:
-        round_up = np.zeros_like(delta, dtype=bool)
-
-    isignificand = np.where(round_up, isignificand + 1, isignificand)
+        should_round_away = ~is_negative & (delta > 0)
+    if rnd == RoundMode.TowardNegative:
+        should_round_away = is_negative & (delta > 0)
+    if rnd == RoundMode.TiesToAway:
+        should_round_away = delta >= 0.5
+    if rnd == RoundMode.TiesToEven:
+        should_round_away = (delta > 0.5) | ((delta == 0.5) & code_is_odd)
+    if rnd == RoundMode.Stochastic:
+        assert srbits is not None
+        should_round_away = delta > (0.5 + srbits) * 2.0**-srnumbits
+    if rnd == RoundMode.StochasticFast:
+        assert srbits is not None
+        should_round_away = delta > srbits * 2.0**-srnumbits
+
+    isignificand = np.where(should_round_away, isignificand + 1, isignificand)
 
     result = np.where(finite_nonzero, np.ldexp(isignificand, expval), absv)
 

src/gfloat/round_ndarray.py

@@ -500,7 +500,10 @@ def test_round_roundtrip(round_float: Callable, fi: FormatInfo) -> None:
         (311, 3, 7.0 / 8),
     ),
 )
-def test_stochastic_rounding(v: float, srnumbits: int, expected_up: float) -> None:
+@pytest.mark.parametrize("impl", ("scalar", "array"))
+def test_stochastic_rounding(
+    impl: bool, v: float, srnumbits: int, expected_up: float
+) -> None:
     fi = format_info_ocp_e5m2
 
     v0 = round_float(fi, v, RoundMode.TowardNegative)
@@ -510,20 +513,26 @@ def test_stochastic_rounding(v: float, srnumbits: int, expected_up: float) -> No
     expected_up_count = expected_up * n
 
     srbits = np.random.randint(0, 2**srnumbits, size=(n,))
-    count_v1 = 0
-    for k in range(n):
-        r = round_float(
-            fi,
-            v,
-            RoundMode.Stochastic,
-            sat=False,
-            srbits=srbits[k],
-            srnumbits=srnumbits,
-        )
-        if r == v1:
-            count_v1 += 1
-        else:
-            assert r == v0
+    if impl == "scalar":
+        count_v1 = 0
+        for k in range(n):
+            r = round_float(
+                fi,
+                v,
+                RoundMode.Stochastic,
+                sat=False,
+                srbits=srbits[k],
+                srnumbits=srnumbits,
+            )
+            if r == v1:
+                count_v1 += 1
+            else:
+                assert r == v0
+    else:
+        vs = np.full(n, v)
+        rs = round_ndarray(fi, vs, RoundMode.Stochastic, False, srbits, srnumbits)
+        assert np.all((rs == v0) | (rs == v1))
+        count_v1 = np.sum(rs == v1)
 
     print(f"SRBits={srnumbits}, observed = {count_v1}, expected = {expected_up_count} ")
     # e.g. if expected is 1250/10000, want to be within 0.5,1.5