ai.djl.nn.transformer

Class MissingOps

java.lang.Object

ai.djl.nn.transformer.MissingOps

public final class MissingOps
extends java.lang.Object

Operators missing from NDArray that are necessary to implement Bert pretraining.

Method Summary

Modifier and Type	Method and Description
static NDArray	gatherNd(NDArray lookup, NDArray indices) Applies the mxnet gather_nd operator.
static NDArray	oneHot(int depth, NDArray data) Creates a one-hot-encoding from the given data.

Methods inherited from class java.lang.Object

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

Method Detail

gatherNd

public static NDArray gatherNd(NDArray lookup,
                               NDArray indices)

Applies the mxnet gather_nd operator.

Parameters:

lookup - array with the data to look up (e.g. an embedding table)

indices - indices to use for lookup

Returns:

the lookup result

oneHot

public static NDArray oneHot(int depth,
                             NDArray data)

Creates a one-hot-encoding from the given data.

Parameters:

depth - size of each one hot encoding (=size of a dictionary)

data - the data to encode

Returns:

the one hot encoding