wip: numeric comparison tooling

Author: hatemhelal

pyscf_ipu/experimental/integrals.py

@@ -24,13 +24,22 @@
 """
 
 
-def factorial(n: IntN, nmax: int = LMAX) -> IntN:
+def factorial_fori(n: IntN, nmax: int = LMAX) -> IntN:
     def body_fun(i, val):
         return val * jnp.where(i <= n, i, 1)
 
     return lax.fori_loop(1, nmax + 1, body_fun, jnp.ones_like(n))
 
 
+def factorial_gamma(n: IntN) -> IntN:
+    from jax.scipy.special import gammaln
+
+    return jnp.exp(gammaln(n + 1))
+
+
+factorial = factorial_fori
+
+
 def factorial2(n: IntN, nmax: int = 2 * LMAX) -> IntN:
     def body_fun(i, val):
         return val * jnp.where((i <= n) & (n % 2 == i % 2), i, 1)

pyscf_ipu/experimental/numerics.py

@@ -0,0 +1,43 @@
+# Copyright (c) 2023 Graphcore Ltd. All rights reserved.
+from functools import wraps
+
+import jax.numpy as jnp
+import numpy as np
+
+
+def apply_fpcast(v, dtype):
+    if isinstance(v, jnp.ndarray) and v.dtype.kind == "f":
+        return v.astype(dtype)
+
+    return v
+
+
+def fpcast(func, dtype=jnp.float32):
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        inputs = [apply_fpcast(v, dtype) for v in args]
+        outputs = func(*inputs, **kwargs)
+        return outputs
+
+    return wrapper
+
+
+def compare_fp32_to_fp64(func):
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        outputs_fp32 = fpcast(func, dtype=jnp.float32)(*args, **kwargs)
+        outputs_fp64 = fpcast(func, dtype=jnp.float64)(*args, **kwargs)
+        print_compare(outputs_fp32, outputs_fp64)
+        return outputs_fp32
+
+    return wrapper
+
+
+def print_compare(fp32, fp64):
+    fp32 = [fp32] if isinstance(fp32, jnp.ndarray) else fp32
+    fp64 = [fp64] if isinstance(fp64, jnp.ndarray) else fp64
+
+    for low, high in zip(fp32, fp64):
+        low = np.asarray(low).astype(np.float64)
+        high = np.asarray(high)
+        print(f" max |fp64 - fp32| = {np.abs(high - low).max()}")

test/test_numerics.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2023 Graphcore Ltd. All rights reserved.
+import jax.numpy as jnp
+from numpy.testing import assert_allclose
+
+from pyscf_ipu.experimental.integrals import factorial_fori, factorial_gamma
+from pyscf_ipu.experimental.numerics import compare_fp32_to_fp64
+
+
+def test_factorial():
+    x = jnp.arange(8, dtype=jnp.float32)
+    y_fori = compare_fp32_to_fp64(factorial_fori)(x, 8)
+    y_gamma = compare_fp32_to_fp64(factorial_gamma)(x)
+    assert_allclose(y_fori, y_gamma, atol=1e-2)