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

By default check_rep is false for shard_map

Browse source
This commit is contained in:
Wassim KABALAN 2024-10-22 17:30:56 -04:00
parent 0bb992fc56
commit 45b2c7f0d2
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/distributed.py
View file

@ -19,7 +19,7 @@ def autoshmap(
gpu_mesh: Mesh | None, gpu_mesh: Mesh | None,
in_specs: Specs, in_specs: Specs,
out_specs: Specs, out_specs: Specs,
check_rep: bool = True, check_rep: bool = False,
auto: frozenset[AxisName] = frozenset()) -> Callable: auto: frozenset[AxisName] = frozenset()) -> Callable:
"""Helper function to wrap the provided function in a shard map if """Helper function to wrap the provided function in a shard map if
the code is being executed in a mesh context.""" the code is being executed in a mesh context."""