mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
3e03deab6f
This makes Placement data representation available in C++ via pybind11. Pull Request resolved: https://github.com/pytorch/pytorch/pull/163030 Approved by: https://github.com/ezyang
21 lines
581 B
Python
21 lines
581 B
Python
# This module is defined in torch/csrc/distributed/python_placement.cpp
|
|
|
|
class Placement:
|
|
def is_partial(self, reduce_op: str | None = None) -> bool: ...
|
|
def is_replicate(self) -> bool: ...
|
|
def is_shard(self, dim: int | None = None) -> bool: ...
|
|
|
|
class Shard(Placement):
|
|
dim: int
|
|
def __init__(self, dim: int): ...
|
|
|
|
class StridedShard(Shard):
|
|
split_factor: int
|
|
def __init__(self, dim: int, *, split_factor: int): ...
|
|
|
|
class Replicate(Placement): ...
|
|
|
|
class Partial(Placement):
|
|
reduce_op: str
|
|
def __init__(self, reduce_op: str | None = None): ...
|