guilhem_lavaux/JaxPM
0
0
Fork 1
mirror of https://github.com/DifferentiableUniverseInitiative/JaxPM.git synced 2025-04-11 21:50:55 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix deprecated FftType in jaxpm.kernels

Browse source
This commit is contained in:
Wassim Kabalan 2024-12-08 22:54:52 +01:00
parent 5d34d3c3a8
commit adaf7d236d
1 changed files with 1 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
jaxpm/kernels.py
View file

@ -1,6 +1,6 @@
import jax.numpy as jnp import jax.numpy as jnp
import numpy as np import numpy as np
from jax.lib.xla_client import FftType from jax.lax import FftType
from jax.sharding import PartitionSpec as P from jax.sharding import PartitionSpec as P
from jaxdecomp import fftfreq3d, get_output_specs from jaxdecomp import fftfreq3d, get_output_specs