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

apply formatting

Browse source
This commit is contained in:
Wassim KABALAN 2024-10-27 00:52:14 +02:00
parent 11f7e90066
commit 4342279817
5 changed files with 26 additions and 22 deletions
Download patch file Download diff file Expand all files Collapse all files

2
jaxpm/distributed.py
View file

@ -40,7 +40,7 @@ def ifft3d(x):
def get_halo_size(halo_size, sharding):
gpu_mesh = sharding.mesh if sharding is not None else None
if gpu_mesh is None or gpu_mesh.empty:
zero_ext = (0, 0, 0)
zero_ext = (0, 0)
zero_tuple = (0, 0)
return (zero_tuple, zero_tuple, zero_tuple), zero_ext
else: