"""
This type stub file was generated by pyright.
"""

import os
import sys
from .dist_graph import DistGraph, DistGraphServer, edge_split, node_split
from .dist_tensor import DistTensor
from .partition import load_partition, load_partition_book, partition_graph
from .graph_partition_book import GraphPartitionBook, PartitionPolicy
from .nn import *
from . import optim
from .rpc import *
from .rpc_server import start_server
from .rpc_client import connect_to_server
from .dist_context import exit_client, initialize
from .kvstore import KVClient, KVServer
from .server_state import ServerState
from .dist_dataloader import DistDataLoader
from .graph_services import in_subgraph, sample_neighbors

"""DGL distributed module contains classes and functions to support
distributed graph neural network training and inference in a cluster of
machines.

This includes a few submodules:

* distributed data structures including distributed graph, distributed tensor
  and distributed embeddings.
* distributed sampling.
* distributed workload split at runtime.
* graph partition.

"""
