Merge pull request #33 from graphcore-research/fix-overflows

Avoid overflow/conversion errors in vectorized code.

Author: awf

pyproject.toml

@@ -41,4 +41,4 @@ ignore_missing_imports = true
 
 [tool.pytest.ini_options]
 addopts = "--nbval"
-testpaths = ["test", "doc"]
+testpaths = ["docs", "test"]

src/gfloat/decode_ndarray.py

@@ -48,18 +48,14 @@ def decode_ndarray(
 
     expBias = fi.expBias
 
-    iszero = (exp == 0) & (significand == 0) & fi.has_zero
-    issubnormal = (exp == 0) & (significand != 0) & fi.has_subnormals
-    isnormal = ~iszero & ~issubnormal
-    expval = np.where(~isnormal, 1 - expBias, exp - expBias)
-    fsignificand = np.where(~isnormal, significand * 2**-t, 1.0 + significand * 2**-t)
-
-    # Normal/Subnormal/Zero case, other values will be overwritten
-    fval = np.where(iszero, 0.0, sign * fsignificand * 2.0**expval)
+    fval = np.zeros_like(codes, dtype=np.float64)
+    isspecial = np.zeros_like(codes, dtype=bool)
 
     if fi.has_infs:
         fval = np.where(codes == fi.code_of_posinf, np.inf, fval)
+        isspecial |= codes == fi.code_of_posinf
         fval = np.where(codes == fi.code_of_neginf, -np.inf, fval)
+        isspecial |= codes == fi.code_of_neginf
 
     if fi.num_nans > 0:
         code_is_nan = codes == fi.code_of_nan
@@ -70,9 +66,21 @@ def decode_ndarray(
             code_is_nan |= abse & (significand >= min_code_with_nan)
 
         fval = np.where(code_is_nan, np.nan, fval)
+        isspecial |= code_is_nan
 
-    # Negative zero
+    # Zero
+    iszero = ~isspecial & (exp == 0) & (significand == 0) & fi.has_zero
+    fval = np.where(iszero, 0.0, fval)
     if fi.has_nz:
         fval = np.where(iszero & (sign < 0), -0.0, fval)
 
+    issubnormal = (exp == 0) & (significand != 0) & fi.has_subnormals
+    expval = np.where(issubnormal, 1 - expBias, exp - expBias)
+    fsignificand = np.where(issubnormal, 0.0, 1.0) + np.ldexp(significand, -t)
+
+    # Normal/Subnormal/Zero case, other values will be overwritten
+    expval_safe = np.where(isspecial | iszero, 0, expval)
+    fval_finite_safe = sign * np.ldexp(fsignificand, expval_safe)
+    fval = np.where(~(iszero | isspecial), fval_finite_safe, fval)
+
     return fval

src/gfloat/round_ndarray.py

@@ -92,7 +92,7 @@ def round_ndarray(
         expval += round_up & (isignificand == 1)
         isignificand = np.where(round_up, 1, isignificand)
 
-    result = np.where(finite_nonzero, isignificand * (2.0**expval), absv)
+    result = np.where(finite_nonzero, np.ldexp(isignificand, expval), absv)
 
     amax = np.where(is_negative, -fi.min, fi.max)
 
@@ -189,10 +189,11 @@ def encode_ndarray(fi: FormatInfo, v: np.ndarray) -> np.ndarray:
         biased_exp = exp.astype(np.int64) + (fi.expBias - 1)
         subnormal_mask = (biased_exp < 1) & fi.has_subnormals
 
-        tsig = np.where(subnormal_mask, sig * 2.0**biased_exp, sig * 2 - 1.0)
+        biased_exp_safe = np.where(subnormal_mask, biased_exp, 0)
+        tsig = np.where(subnormal_mask, np.ldexp(sig, biased_exp_safe), sig * 2 - 1.0)
         biased_exp[subnormal_mask] = 0
 
-        isig = np.floor(tsig * 2**t).astype(int)
+        isig = np.floor(np.ldexp(tsig, t)).astype(np.int64)
 
         zero_mask = fi.has_zero & (isig == 0) & (biased_exp == 0)
         if not fi.has_nz:

test/test_decode.py

@@ -260,7 +260,10 @@ def test_consistent_decodes_all_values(
     npivals = np.arange(
         np.iinfo(int_dtype).min, int(np.iinfo(int_dtype).max) + 1, dtype=int_dtype
     )
-    npfvals = npivals.view(dtype=npfmt).astype(np.float64)
+
+    with np.errstate(invalid="ignore"):
+        # Warning here when converting bfloat16 NaNs to float64
+        npfvals = npivals.view(dtype=npfmt).astype(np.float64)
 
     # Scalar version
     for i, npfval in zip(npivals, npfvals):