Package ai.djl.nn.transformer

Class IdEmbedding

java.lang.Object

ai.djl.nn.AbstractBaseBlock

ai.djl.nn.AbstractBlock

ai.djl.nn.transformer.IdEmbedding

All Implemented Interfaces:

Block

public final class IdEmbedding
extends AbstractBlock

An Embedding from integer ids to float vectors. Output shape is the input shape + one dimension for the embedding. E.g. If input shape is (-1, 128), embedding size is 1024, then the output shape is (-1, 128, 1024)

Nested Class Summary

Nested Classes

Modifier and Type	Class	Description
static final class	IdEmbedding.Builder	The Builder to construct an IdEmbedding type of Block.

Field Summary

Fields inherited from class ai.djl.nn.AbstractBlock

children, parameters

Fields inherited from class ai.djl.nn.AbstractBaseBlock

inputNames, inputShapes, outputDataTypes, version

Method Summary

Modifier and Type	Method	Description
protected NDList	forwardInternal(ParameterStore ps, NDList inputs, boolean training, ai.djl.util.PairList<String,Object> params)	A helper for Block.forward(ParameterStore, NDList, boolean, PairList) after initialization.
Shape[]	getOutputShapes(Shape[] inputShapes)	Returns the expected output shapes of the block for the specified input shapes.
NDArray	getValue(ParameterStore ps, Device device, boolean training)	Quick hack for bert model to acces embedding table, replace by a proper function to calculate raw logits from embeddings.
void	initializeChildBlocks(NDManager manager, DataType dataType, Shape... inputShapes)	Initializes the Child blocks of this block.
NDArray	probabilities(ParameterStore parameterStore, NDArray input, boolean training)	Turns an array of embeddings of shape (d0 ...

Methods inherited from class ai.djl.nn.AbstractBlock

addChildBlock, addChildBlock, addChildBlockSingleton, addParameter, getChildren, getDirectParameters

Methods inherited from class ai.djl.nn.AbstractBaseBlock

beforeInitialize, cast, clear, describeInput, forward, forward, forwardInternal, getInputShapes, getOutputDataTypes, getParameters, initialize, isInitialized, loadMetadata, loadParameters, prepare, readInputShapes, saveInputShapes, saveMetadata, saveParameters, setInitializer, setInitializer, setInitializer, toString

Methods inherited from class java.lang.Object

clone, equals, finalize, getClass, hashCode, notify, notifyAll, wait, wait, wait

Methods inherited from interface ai.djl.nn.Block

forward, freezeParameters, freezeParameters, getOutputShapes

Method Details

getOutputShapes

public Shape[] getOutputShapes(Shape[] inputShapes)

Returns the expected output shapes of the block for the specified input shapes.

Parameters:

inputShapes - the shapes of the inputs

Returns:

the expected output shapes of the block

forwardInternal

protected NDList forwardInternal(ParameterStore ps,
 NDList inputs,
 boolean training,
 ai.djl.util.PairList<String,Object> params)

A helper for Block.forward(ParameterStore, NDList, boolean, PairList) after initialization.

Specified by:

forwardInternal in class AbstractBaseBlock

Parameters:

ps - the parameter store

inputs - the input NDList

training - true for a training forward pass

params - optional parameters

Returns:

the output of the forward pass

probabilities

public NDArray probabilities(ParameterStore parameterStore,
 NDArray input,
 boolean training)

Turns an array of embeddings of shape (d0 ... dN, E) into an array of log probabilities of shape (d0 ... dN, D) that shows the probability distribution that a given embedding corresponds to an entry in the internal embedding table.

Parameters:

parameterStore - the parameters store

input - the embeddings to create log probabilities for

training - true for a training forward pass

Returns:

log probabilities for each embedding

getValue

public NDArray getValue(ParameterStore ps,
 Device device,
 boolean training)

Quick hack for bert model to acces embedding table, replace by a proper function to calculate raw logits from embeddings. TODO: replace by function to get logits

Parameters:

ps - the parameter store

device - device to get internal table for

training - true for a training forward pass

Returns:

this embedding table as an array on the given device

initializeChildBlocks

public void initializeChildBlocks(NDManager manager,
 DataType dataType,
 Shape... inputShapes)

Initializes the Child blocks of this block. You need to override this method if your subclass has child blocks. Used to determine the correct input shapes for child blocks based on the requested input shape for this block.

Overrides:

initializeChildBlocks in class AbstractBaseBlock

Parameters:

manager - the manager to use for initialization

dataType - the requested data type

inputShapes - the expected input shapes for this block