MCPcopy Create free account
hub / github.com/NVIDIA/TensorRT-LLM / __init__

Method __init__

tensorrt_llm/layers/normalization.py:98–120  ·  view source on GitHub ↗
(self,
                 num_groups,
                 num_channels,
                 eps=1e-05,
                 affine=True,
                 dtype=None)

Source from the content-addressed store, hash-verified

96class GroupNorm(Module):
97
98 def __init__(self,
99 num_groups,
100 num_channels,
101 eps=1e-05,
102 affine=True,
103 dtype=None):
104 super().__init__()
105
106 if num_channels % num_groups != 0:
107 raise ValueError('num_channels must be divisible by num_groups')
108
109 self.num_groups = num_groups
110 self.num_channels = num_channels
111 self.affine = affine
112
113 if self.affine:
114 self.weight = Parameter(shape=(self.num_channels, ), dtype=dtype)
115 self.bias = Parameter(shape=(self.num_channels, ), dtype=dtype)
116 else:
117 self.register_parameter('weight', None)
118 self.register_parameter('bias', None)
119
120 self.eps = eps
121
122 def forward(self, x):
123 weight = None if self.weight is None else self.weight.value

Callers

nothing calls this directly

Calls 3

ParameterClass · 0.85
__init__Method · 0.45
register_parameterMethod · 0.45

Tested by

no test coverage detected