/
githubmirror
/
ColossalAI
Обзор
Документация
Войти
/
githubmirror
/
ColossalAI
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
main
colossalai/legacy/tensor/tensor_spec.py
21 строка
735 B
Rocky Duan
Fix ColoTensorSpec for py11 (#5440)
26 мар 2024, 10:56
Не верифицирован
26 мар 2024, 10:56
cbe34c5
Код
Авторство
О чём код?
from dataclasses import dataclass, field from typing import Optional from colossalai.legacy.tensor.distspec import DistPlacementPattern, _DistSpec from colossalai.legacy.tensor.process_group import ProcessGroup from .compute_spec import ComputeSpec @dataclass class ColoTensorSpec: """ColoTensorSpec A data class for specifications of the `ColoTensor`. It contains attributes of `ProcessGroup`, `_DistSpec`, `ComputeSpec`. The latter two attributes are optional. If not set, they are default value is `Replicate()` and `None`. """ pg: ProcessGroup dist_attr: Optional[_DistSpec] = field(default_factory=lambda: _DistSpec(DistPlacementPattern.REPLICATE)) compute_attr: Optional[ComputeSpec] = None