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

Method __init__

tensorrt_llm/layers/normalization.py:28–54  ·  view source on GitHub ↗
(self,
                 normalized_shape,
                 eps=1e-05,
                 elementwise_affine=True,
                 bias=True,
                 dtype=None,
                 tp_size=1,
                 tp_dim=-1)

Source from the content-addressed store, hash-verified

26class LayerNorm(Module):
27
28 def __init__(self,
29 normalized_shape,
30 eps=1e-05,
31 elementwise_affine=True,
32 bias=True,
33 dtype=None,
34 tp_size=1,
35 tp_dim=-1):
36 super().__init__()
37 if isinstance(normalized_shape, int):
38 normalized_shape = (normalized_shape, )
39 self.normalized_shape = tuple(normalized_shape)
40 self.elementwise_affine = elementwise_affine
41 if self.elementwise_affine:
42 self.weight = Parameter(shape=self.normalized_shape, dtype=dtype)
43 if bias:
44 self.bias = Parameter(shape=self.normalized_shape, dtype=dtype)
45 else:
46 self.register_parameter('bias', None)
47 else:
48 self.register_parameter('weight', None)
49 self.register_parameter('bias', None)
50
51 self.eps = eps
52 self.dtype = dtype
53 self.tp_size = tp_size
54 self.tp_dim = tp_dim
55
56 def forward(self, x, normalized_shape=None):
57 weight = 1. 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