guilhem_lavaux/JaxPM
0
0
Fork 1
mirror of https://github.com/DifferentiableUniverseInitiative/JaxPM.git synced 2025-06-28 08:01:12 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

?

Browse source 
clean
This commit is contained in:
chenze@igpu 2025-06-10 22:46:26 +09:00
parent cb2a7ab17f
commit cf979f3250
1 changed files with 16 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

17
jaxpm/distributed.py
View file

@ -79,6 +79,21 @@ def slice_unpad_impl(x, pad_width):
return x[tuple(unpad_slice)]
def slice_pad_impl(x, pad_width):
x = jnp.pad(x, pad_width=pad_width)
halo_x, _ = pad_width[0]
halo_y, _ = pad_width[1]
# Apply corrections along x
x = x.at[halo_x // 2:halo_x].add(x[halo_x:halo_x + halo_x // 2])
x = x.at[-halo_x:-halo_x // 2].add(x[-(halo_x + halo_x // 2):-halo_x])
# Apply corrections along y
x = x.at[:, halo_y // 2:halo_y].add(x[:, halo_y:halo_y + halo_y // 2])
x = x.at[:, -halo_y:-halo_y // 2].add(x[:, -(halo_y + halo_y // 2):-halo_y])
return x
def slice_pad(x, pad_width, sharding):
gpu_mesh = sharding.mesh if sharding is not None else None
@ -86,7 +101,7 @@ def slice_pad(x, pad_width, sharding):
pad_width[0][0] > 0 or pad_width[1][0] > 0):
assert sharding is not None
spec = sharding.spec
return shard_map((partial(jnp.pad, pad_width=pad_width)),
return shard_map((partial(slice_pad_impl, pad_width=pad_width)),
mesh=gpu_mesh,
in_specs=spec,
out_specs=spec)(x)