change jax version and try different tf

Hazboun6 · Dec 31, 2024 · 6e53242 · 6e53242
1 parent d093f1d
commit 6e53242
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/hasasia/sensitivity.py b/hasasia/sensitivity.py
@@ -169,8 +169,8 @@ def get_Tf(designmatrix, toas, N=None, nf=200, fmin=None, fmax=2e-7,
         else:
             G = Gmatrix
         m = G.shape[1]
-        Gtilde = jnp.zeros((ff.size, G.shape[1]), dtype=jnp.complex128)
-        Gtilde = np.dot(np.exp(1j*2*np.pi*ff[:,np.newaxis]*toas),G).astype(jnp.complex128)
+        Gtilde = jnp.zeros((ff.size, G.shape[1]), dtype=jnp.complex64)
+        Gtilde = np.dot(np.exp(1j*2*np.pi*ff[:,np.newaxis]*toas),G).astype(jnp.complex64)
         Tmat = jnp.matmul(jnp.conjugate(Gtilde),Gtilde.T)/N_TOA
         if twofreqs:
             Tmat = np.real(Tmat)

diff --git a/requirements.txt b/requirements.txt
@@ -1,7 +1,7 @@
 numpy>=1.14
 scipy>=1.0
 astropy>=3.0
-jaxlib
-jax
-healpy
+jaxlib>=0.3.0
+jax[cpu]~=0.4
+healpy>=1.14