guilhem_lavaux/JaxPM
0
0
Fork 1
mirror of https://github.com/DifferentiableUniverseInitiative/JaxPM.git synced 2025-06-29 16:41:11 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

merge with JZ

Browse source
This commit is contained in:
Wassim KABALAN 2024-10-18 14:59:40 -04:00
commit ab86699c88
7 changed files with 226 additions and 131 deletions
Download patch file Download diff file Expand all files Collapse all files

4
jaxpm/distributed.py
View file

@ -91,8 +91,8 @@ def get_halo_size(halo_size):
def halo_exchange(x, halo_extents, halo_periods=(True, True, True)):
mesh = mesh_lib.thread_resources.env.physical_mesh
if distributed and not (mesh.empty) and (halo_extents[0] > 0
or halo_extents[1] > 0):
if distributed and not (mesh.empty) and (halo_extents > 0
or halo_extents > 0):
return jaxdecomp.halo_exchange(x, halo_extents, halo_periods)
else:
return x