text-generation-inference/server/text_generation_server/layers/__init__.py

from text_generation_server.layers.tensor_parallel import (
    TensorParallelColumnLinear,
    TensorParallelRowLinear,
    TensorParallelEmbedding,
)
from text_generation_server.layers.linear import (
    get_linear,
    FastLinear,
)
from text_generation_server.layers.speculative import SpeculativeHead

# Just to add the `load` methods.
from text_generation_server.layers.layernorm import load_layer_norm
from text_generation_server.layers.conv import load_conv2d