from __future__ import annotations

from collections import defaultdict as defaultdict
from typing import Any, Optional

from .dist_attribute import OperatorDistributedAttribute as OperatorDistributedAttribute
from .dist_attribute import TensorDistributedAttribute as TensorDistributedAttribute
from .dist_attribute import append_op_input_suffix as append_op_input_suffix
from .dist_attribute import append_op_output_suffix as append_op_output_suffix
from .dist_attribute import get_op_dist_attr_field_keys as get_op_dist_attr_field_keys
from .dist_attribute import (
    get_tensor_dist_attr_field_keys as get_tensor_dist_attr_field_keys,
)

class DistributedOperator:
    def __init__(self, serial_op: Any, dist_attr: Optional[Any] = ...) -> None: ...
    @property
    def serial_op(self): ...
    @property
    def dist_attr(self): ...
    @dist_attr.setter
    def dist_attr(self, dist_attr: Any) -> None: ...
    def get_serial_input(self, name: Any): ...
    def get_serial_output(self, name: Any): ...
    def validate_dist_attr(self): ...
    def __deepcopy__(self, memo: Any): ...

class DistributedModule:
    def __init__(self, serial_module: Any, dist_attr: Optional[Any] = ...) -> None: ...
    def __call__(self, *args: Any, **kwargs: Any): ...
