tf.EmbeddingLayer module

tf.EmbeddingLayer module

class tf.EmbeddingLayer.EmbeddingLayer(*args: Any, **kwargs: Any)

Bases: modelzoo.common.layers.tf.BaseLayer.BaseLayer

Embedding layer. Built on top of the Keras Embedding layer.

build(input_shape)
call(inputs, pad_id=- 1, scale=1)

Get token embeddings of inputs.

Parameters

  • inputs (Tensor) – A tensor with shape [batch_size, length].

  • pad_id – Integer specifying which input ID corresponds instead to padding. It does not need to be a legal vocabulary entry. Any `inputs` elements equal to this value will not be looked up, but instead directly output zeros. On the Wafer Scale Engine, this indicates the presence of variable sequence length.

  • scale – Scaling of the embedding (in MLPERF hidden_size**0.5 is used).

Returns

A tensor of embeddings with shape [batch_size, length, hidden_size]. Padded positions are filled with zeros.

Return type

embeddings (Tensor)

embedding_table()

previous

tf.DropoutLayer module

next

tf.FeedForwardNetwork module