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

quick fix in distributed

Browse source
This commit is contained in:
Wassim KABALAN 2024-10-22 12:15:37 -04:00
parent a8b194f326
commit 0433c615f3
1 changed files with 3 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
jaxpm/distributed.py
View file

@ -153,9 +153,9 @@ def normal_field(mesh_shape, seed, sharding):
def normal(keys, shape, dtype):
idx = lax.axis_index(x_axis)
if not single_axis:
y_index = lax.axis_index(y_axis)
x_size = lax.psum(1, axis_name=x_axis)
idx += y_index * x_size
y_index = lax.axis_index(y_axis)
x_size = lax.psum(1, axis_name=x_axis)
idx += y_index * x_size
return jax.random.normal(key=keys[idx], shape=shape, dtype=dtype)