fix deprecated FftType in jaxpm.kernels

DifferentiableUniverseInitiative · Dec 8, 2024 · adaf7d2 · adaf7d2
1 parent 5d34d3c
commit adaf7d2
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/jaxpm/kernels.py b/jaxpm/kernels.py
@@ -1,6 +1,6 @@
 import jax.numpy as jnp
 import numpy as np
-from jax.lib.xla_client import FftType
+from jax.lax import FftType
 from jax.sharding import PartitionSpec as P
 from jaxdecomp import fftfreq3d, get_output_specs