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

Method __init__

tensorrt_llm/layers/embedding.py:574–599  ·  view source on GitHub ↗
(self,
                 in_features,
                 hidden_size,
                 out_features=None,
                 act_fn="gelu_tanh",
                 mapping=None,
                 dtype=None)

Source from the content-addressed store, hash-verified

572 """
573
574 def __init__(self,
575 in_features,
576 hidden_size,
577 out_features=None,
578 act_fn="gelu_tanh",
579 mapping=None,
580 dtype=None):
581 super().__init__()
582 if out_features is None:
583 out_features = hidden_size
584 tp_group = mapping.tp_group
585 tp_size = mapping.tp_size
586 self.linear_1 = ColumnLinear(in_features=in_features,
587 out_features=hidden_size,
588 bias=True,
589 tp_group=tp_group,
590 tp_size=tp_size,
591 dtype=dtype,
592 gather_output=False)
593 self.act_1 = ACT2FN[act_fn]
594 self.linear_2 = RowLinear(in_features=hidden_size,
595 out_features=out_features,
596 bias=True,
597 tp_group=tp_group,
598 tp_size=tp_size,
599 dtype=dtype)
600
601 def forward(self, caption):
602 hidden_states = self.linear_1(caption)

Callers

nothing calls this directly

Calls 2

RowLinearClass · 0.85
__init__Method · 0.45

Tested by

no test coverage detected