tensorplay.distributed.tensor.distribute_module

tensorplay.distributed.tensor.distribute_module(module: Any, device_mesh: DeviceMesh | None = None, partition_fn: Callable[[str, Any, DeviceMesh], Any] | None = None, input_fn: Callable[[Any, Any, DeviceMesh], Any] | None = None, output_fn: Callable[[Any, Any, DeviceMesh], Any] | None = None) → Any[source]