mirror of
https://github.com/DifferentiableUniverseInitiative/JaxPM.git
synced 2025-04-07 20:30:54 +00:00
By default check_rep is false for shard_map
This commit is contained in:
parent
0bb992fc56
commit
45b2c7f0d2
1 changed files with 1 additions and 1 deletions
|
@ -19,7 +19,7 @@ def autoshmap(
|
|||
gpu_mesh: Mesh | None,
|
||||
in_specs: Specs,
|
||||
out_specs: Specs,
|
||||
check_rep: bool = True,
|
||||
check_rep: bool = False,
|
||||
auto: frozenset[AxisName] = frozenset()) -> Callable:
|
||||
"""Helper function to wrap the provided function in a shard map if
|
||||
the code is being executed in a mesh context."""
|
||||
|
|
Loading…
Add table
Reference in a new issue