By default check_rep is false for shard_map

This commit is contained in:
Wassim KABALAN 2024-10-22 17:30:56 -04:00
parent 0bb992fc56
commit 45b2c7f0d2

View file

@ -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."""