dtensor.run_on() only accept Mesh
Removal of layout_or: the result does not ensure the tensors follow layout,
passing in a Layout is confusing.
Existing users of layout_or path can migrate to to:
dtensor.run_on(layout.mesh)
Also, to use the context conditionally:
with dtensor.run_on(mesh) if mesh else contextlib.nullcontext():
@@ -157,7 +157,7 @@ def call_with_layout(fn, layout, *args, **kwargs):
The output of fn, with potential relayout with the layout specified.
"""
if layout:
- with dtensor.run_on(layout):
+ with dtensor.run_on(layout.mesh):
result = fn(*args, **kwargs)
return dtensor.relayout(result, layout)
return fn(*args, **kwargs)
