tf.layers.EmbeddingLayer module#

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

Bases: modelzoo.common.tf.layers.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()#