phivenv/Lib/site-packages/torch/distributed/pipelining/__init__.py

# Copyright (c) Meta Platforms, Inc. and affiliates
from ._IR import Pipe, pipe_split, pipeline, SplitPoint
from .schedules import (
    _ScheduleForwardOnly,
    Schedule1F1B,
    ScheduleGPipe,
    ScheduleInterleaved1F1B,
    ScheduleInterleavedZeroBubble,
    ScheduleLoopedBFS,
    ScheduleZBVZeroBubble,
)
from .stage import build_stage, PipelineStage


__all__ = [
    "Pipe",
    "pipe_split",
    "SplitPoint",
    "pipeline",
    "PipelineStage",
    "build_stage",
    "Schedule1F1B",
    "ScheduleGPipe",
    "ScheduleInterleaved1F1B",
    "ScheduleLoopedBFS",
    "ScheduleInterleavedZeroBubble",
    "ScheduleZBVZeroBubble",
]