From 8e8e8964befeadb12c12d2287e250f02baddfaff Mon Sep 17 00:00:00 2001
From: Wassim KABALAN <wastondev@gmail.com>
Date: Fri, 25 Oct 2024 10:15:43 +0200
Subject: [PATCH] jaxDecomp pfft3d promotes to complex automatically

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

diff --git a/jaxpm/distributed.py b/jaxpm/distributed.py
index ab85856..4fdb764 100644
--- a/jaxpm/distributed.py
+++ b/jaxpm/distributed.py
@@ -30,7 +30,7 @@ def autoshmap(
 
 
 def fft3d(x):
-    return jaxdecomp.pfft3d(x.astype(jnp.complex64))
+    return jaxdecomp.pfft3d(x)
 
 
 def ifft3d(x):