modular/max/nn/__init__.py at main · avioligo/modular · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
# ===----------------------------------------------------------------------=== #
# Copyright (c) 2025, Modular Inc. All rights reserved.
#
# Licensed under the Apache License v2.0 with LLVM Exceptions:
# https://llvm.org/LICENSE.txt
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ===----------------------------------------------------------------------=== #

from .attention import (
    AttentionImpl,
    AttentionImplQKV,
    AttentionWithRope,
    AttentionWithRopeQKV,
    AttentionWithRopeV1,
    DistributedAttentionImpl,
    DistributedAttentionWithRope,
    GGUFQAttentionWithRope,
    GPTQAttentionWithRope,
    RaggedAttention,
)
from .clamp import clamp
from .comm import Allreduce, Signals
from .conv import Conv1D, Conv1DV1, Conv2D, Conv2DV1, Conv3D, Conv3DV1
from .conv_transpose import ConvTranspose1d, WeightNormConvTranspose1d
from .embedding import Embedding, EmbeddingV1, VocabParallelEmbedding
from .layer import Layer, LayerList, Module
from .linear import (
    MLP,
    MLPV1,
    ColumnParallelLinear,
    DistributedMLP,
    Float8Config,
    Float8InputScaleSpec,
    Float8ScaleGranularity,
    Float8ScaleOrigin,
    Float8WeightScaleSpec,
    GPTQLinear,
    Linear,
    LinearV1,
)
from .norm import (
    DistributedRMSNorm,
    GroupNorm,
    LayerNorm,
    LayerNormV1,
    RMSNorm,
    RMSNormV1,
)
from .rotary_embedding import (
    LinearScalingParams,
    Llama3RopeScalingParams,
    Llama3RotaryEmbedding,
    OptimizedRotaryEmbedding,
    RotaryEmbedding,
)
from .sequential import Sequential
from .transformer import (
    DistributedTransformer,
    DistributedTransformerBlock,
    ReturnLogits,
    Transformer,
    TransformerBlock,
)

__all__ = [
    "Allreduce",
    "AttentionImpl",
    "AttentionImplQKV",
    "AttentionWithRopeV1",
    "AttentionWithRopeQKV",
    "AttentionWithRope",
    "RaggedAttention",
    "clamp",
    "Conv1DV1",
    "Conv2DV1",
    "Conv3DV1",
    "Conv1D",
    "Conv2D",
    "Conv3D",
    "ConvTranspose1d",
    "WeightNormConvTranspose1d",
    "DistributedAttentionImpl",
    "DistributedAttentionWithRope",
    "ColumnParallelLinear",
    "DistributedMLP",
    "DistributedRMSNorm",
    "DistributedTransformer",
    "DistributedTransformerBlock",
    "EmbeddingV1",
    "Embedding",
    "Float8Config",
    "Float8ScaleGranularity",
    "Float8ScaleOrigin",
    "Float8InputScaleSpec",
    "Float8WeightScaleSpec",
    "GGUFQAttentionWithRope",
    "GPTQAttentionWithRope",
    "GPTQLinear",
    "GroupNorm",
    "Layer",
    "LayerList",
    "LayerNormV1",
    "LayerNorm",
    "LinearV1",
    "Linear",
    "LinearScalingParams",
    "Llama3RopeScalingParams",
    "Llama3RotaryEmbedding",
    "MLPV1",
    "MLP",
    "Module",
    "OptimizedRotaryEmbedding",
    "RMSNormV1",
    "RMSNorm",
    "RotaryEmbedding",
    "ReturnLogits",
    "Sequential",
    "Signals",
    "Transformer",
    "TransformerBlock",
    "VocabParallelEmbedding",
]