use jnp in primitive product method to support jit

graphcore-research · Oct 15, 2023 · ce026e9 · ce026e9
1 parent 04a9de3
commit ce026e9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pyscf_ipu/experimental/primitive.py b/pyscf_ipu/experimental/primitive.py
@@ -41,7 +41,7 @@ def product(a: Primitive, b: Primitive) -> Primitive:
     lmn = a.lmn + b.lmn
     c = a.norm * b.norm
     Rab = a.center - b.center
-    c *= np.exp(-a.alpha * b.alpha / alpha * np.inner(Rab, Rab))
+    c *= jnp.exp(-a.alpha * b.alpha / alpha * jnp.inner(Rab, Rab))
     return Primitive(center=center, alpha=alpha, lmn=lmn, norm=c)