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

Method __init__

tensorrt_llm/layers/embedding.py:127–140  ·  view source on GitHub ↗
(self,
                 num_embeddings,
                 embedding_dim,
                 vocab_size=None,
                 dtype=None,
                 tp_size=1,
                 tp_group=None,
                 sharding_dim=0,
                 tp_rank=0)

Source from the content-addressed store, hash-verified

125 """
126
127 def __init__(self,
128 num_embeddings,
129 embedding_dim,
130 vocab_size=None,
131 dtype=None,
132 tp_size=1,
133 tp_group=None,
134 sharding_dim=0,
135 tp_rank=0):
136 super().__init__(num_embeddings, embedding_dim, dtype, tp_size,
137 tp_group, sharding_dim, tp_rank)
138 if vocab_size is None:
139 vocab_size = num_embeddings
140 self.vocab_size = vocab_size
141
142 def forward(self, tokens, prompt_embedding_table, tasks, task_vocab_size):
143 """

Callers

nothing calls this directly

Calls 1

__init__Method · 0.45

Tested by

no test coverage detected