wehub-resource-sync
44 lines
1.3 KiB
Python
44 lines
1.3 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
from ...inference_utils import ActivationType, DtypeEnum
|
|
from ...modules.ds_module import DSModuleConfig
|
|
|
|
|
|
class DSLinearConfig(DSModuleConfig):
|
|
"""
|
|
Config class for DSLinearBase.
|
|
"""
|
|
|
|
in_channels: int
|
|
"""
|
|
Number of input channels
|
|
"""
|
|
|
|
out_channels: int
|
|
"""
|
|
Number of output channels. NOTE: If this linear layer is using a gated activation function,
|
|
the value for ``out_channels`` passed here should refer to the number of channels after
|
|
gating (i.e., the expected weight shape before transformations will be ``[out_channels * 2, in_channels]``).
|
|
"""
|
|
|
|
activation: ActivationType = ActivationType.IDENTITY
|
|
"""
|
|
The activation function for this layer. See :class:`deepspeed.inference.inference_utils.ActivationType` for
|
|
supported activation functions.
|
|
"""
|
|
|
|
input_dtype: DtypeEnum = DtypeEnum.fp16
|
|
"""
|
|
The data type of the input tensor. See :class:`deepspeed.inference.inference_utils.DtypeEnum` for supported
|
|
data types.
|
|
"""
|
|
|
|
output_dtype: DtypeEnum = DtypeEnum.fp16
|
|
"""
|
|
The data type of the output tensor. See :class:`deepspeed.inference.inference_utils.DtypeEnum` for supported
|
|
data types.
|
|
"""
|