From 23a93deec8c783a9e45a4ecc1b689f031632ed00 Mon Sep 17 00:00:00 2001
From: Yin Li <eelregit@gmail.com>
Date: Fri, 10 Mar 2023 22:03:41 +0800
Subject: [PATCH] Change convolution kernel dtype for JAX backend

---
 mcfit/mcfit.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mcfit/mcfit.py b/mcfit/mcfit.py
index f848a4c..9fa24a5 100644
--- a/mcfit/mcfit.py
+++ b/mcfit/mcfit.py
@@ -195,7 +195,7 @@ def _setup(self):
         m = numpy.arange(0, self.N//2 + 1)
         self._u = self.MK(self.q + 2j * math.pi / self.N / Delta * m)
         self._u *= numpy.exp(-2j * math.pi * lnxy / self.N / Delta * m)
-        self._u = self.np.asarray(self._u)
+        self._u = self.np.asarray(self._u, dtype=(self.x[0] + 0j).dtype)
 
         # following is unnecessary because hfft ignores the imag at Nyquist anyway
         #if not self.lowring and self.N % 2 == 0: