You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
22 lines
598 B
22 lines
598 B
# 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(Placement):
|
|
dim: int
|
|
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): ...
|