jaxdecomp 0.2.2__tar.gz → 0.2.3__tar.gz

{jaxdecomp-0.2.2 → jaxdecomp-0.2.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxdecomp
-Version: 0.2.2
+Version: 0.2.3
 Summary: JAX bindings for the cuDecomp library
 Author: Wassim Kabalan, Francois Lanusse
 License: MIT License
@@ -34,9 +34,9 @@ Classifier: Intended Audience :: Developers
 Classifier: Intended Audience :: Science/Research
 Project-URL: Homepage, https://github.com/DifferentiableUniverseInitiative/jaxDecomp
 Requires-Dist: jaxtyping>=0.2.0
-Requires-Dist: jax>=0.4.30
+Requires-Dist: jax>=0.4.35
 Requires-Dist: pytest>=8.0.0; extra == "test"
-Requires-Dist: jax>=0.4.30; extra == "test"
+Requires-Dist: jax>=0.4.35; extra == "test"
 Provides-Extra: test
 Description-Content-Type: text/markdown
 

{jaxdecomp-0.2.2 → jaxdecomp-0.2.3}/pyproject.toml

@@ -3,7 +3,7 @@ requires = ["scikit-build-core>=0.4.0", "pybind11>=2.9.0"]
 build-backend = "scikit_build_core.build"
 [project]
 name = "jaxdecomp"
-version = "0.2.2"
+version = "0.2.3"
 description = "JAX bindings for the cuDecomp library"
 authors = [
   { name = "Wassim Kabalan" },
@@ -24,11 +24,11 @@ classifiers = [
 ]
 dependencies = [
 "jaxtyping>=0.2.0",
-"jax>=0.4.30",
+"jax>=0.4.35",
 ]
 
 [project.optional-dependencies]
-test = ["pytest>=8.0.0" , "jax>=0.4.30"]
+test = ["pytest>=8.0.0" , "jax>=0.4.35"]
 
 [tool.scikit-build]
 minimum-version = "0.8"

{jaxdecomp-0.2.2 → jaxdecomp-0.2.3}/src/jaxdecomp/fft.py

@@ -2,7 +2,8 @@ from functools import partial
 from typing import Optional, Sequence
 
 import jax.numpy as jnp
-from jax import jit
+from jax import jit, lax
+from jax._src import dtypes
 from jax._src.typing import Array, ArrayLike
 
 from jaxdecomp._src.cudecomp.fft import pfft as _cudecomp_pfft
@@ -115,6 +116,14 @@ def _do_pfft(func_name: str,
     typ = fft_type
   else:
     raise TypeError(f"Unknown FFT type value '{fft_type}'")
+
+  match typ:
+    case FftType.FFT | FftType.IFFT:
+      arr = lax.convert_element_type(arr,
+                                     dtypes.to_complex_dtype(dtypes.dtype(arr)))
+    case FftType.RFFT | FftType.IRFFT:
+      raise ValueError("Not implemented wait (SOON)")
+
   if backend.lower() == "cudecomp":
     transformed = _cudecomp_pfft(arr, typ)
   elif backend.lower() == "jax":