BUG: add vectorize

JAXtronomy · Oct 18, 2024 · 82617f0 · 82617f0
1 parent d6e975c
commit 82617f0
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/spexial/_src/spence.py b/src/spexial/_src/spence.py
@@ -13,6 +13,7 @@
 MAXITER = 500
 
 
+@jax.numpy.vectorize
 def complex_spence_series_0(z):
     """
     Small z branch, uses a series expansion about :math:`z = 0`
@@ -27,6 +28,7 @@ def complex_spence_series_0(z):
     return np.pi**2 / 6 - sum1 + jnp.log(z) * sum2
 
 
+@jax.numpy.vectorize
 def complex_spence_series_1(z):
     """
     Middle branch, an expansion around :math:`z = 1`