Package ai.djl.nn.transformer

Class MissingOps

  • public final class MissingOps
extends java.lang.Object
    Operators missing from NDArray that are necessary to implement Bert pretraining.

    • Method Summary

      All Methods Static Methods Concrete Methods 
      Modifier and Type Method 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