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.

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