ai.djl.mxnet.jna

Class JnaUtils

java.lang.Object

ai.djl.mxnet.jna.JnaUtils

public final class JnaUtils
extends java.lang.Object

A class containing utilities to interact with the MXNet Engine's Java Native Access (JNA) layer.

Nested Class Summary

Nested Classes

Modifier and Type	Class and Description
static class	JnaUtils.NumpyMode An enum that enumerates the statuses of numpy mode.

Field Summary

Fields

Modifier and Type	Field and Description
static java.lang.String[]	EMPTY_ARRAY
static java.lang.String	MXNET_THREAD_SAFE_PREDICTOR

Method Summary

Modifier and Type	Method and Description
static void	autogradBackward(ai.djl.ndarray.NDList array, int retainGraph)
static void	autogradBackwardExecute(int numOutput, ai.djl.ndarray.NDList array, ai.djl.ndarray.NDArray outgrad, int numVariables, com.sun.jna.Pointer varHandles, int retainGraph, int createGraph, int isTrain, com.sun.jna.Pointer gradHandles, com.sun.jna.Pointer gradSparseFormat)
static com.sun.jna.Pointer	autogradGetSymbol(ai.djl.ndarray.NDArray array)
static boolean	autogradIsRecording()
static boolean	autogradIsTraining()
static void	autogradMarkVariables(int numVar, com.sun.jna.Pointer varHandles, java.nio.IntBuffer reqsArray, com.sun.jna.Pointer gradHandles)
static boolean	autogradSetIsRecording(boolean isRecording)
static boolean	autogradSetTraining(boolean isTraining)
static MxNDArray[]	cachedOpInvoke(MxNDManager manager, com.sun.jna.Pointer cachedOpHandle, MxNDArray[] inputs)
static void	checkCall(int ret)
static CachedOp	createCachedOp(MxSymbolBlock block, MxNDManager manager) Creates cached op flags.
static com.sun.jna.Pointer	createNdArray(ai.djl.Device device, ai.djl.ndarray.types.Shape shape, ai.djl.ndarray.types.DataType dtype, int size, boolean delayedAlloc)
static com.sun.jna.Pointer	createSparseNdArray(ai.djl.ndarray.types.SparseFormat fmt, ai.djl.Device device, ai.djl.ndarray.types.Shape shape, ai.djl.ndarray.types.DataType dtype, ai.djl.ndarray.types.DataType[] auxDTypes, ai.djl.ndarray.types.Shape[] auxShapes, boolean delayedAlloc)
static com.sun.jna.Pointer	createSymbolFromFile(java.lang.String path)
static void	freeCachedOp(com.sun.jna.Pointer handle)
static void	freeNdArray(com.sun.jna.Pointer ndArray)
static void	freeSymbol(com.sun.jna.Pointer symbol)
static java.util.Set<java.lang.String>	getAllOpNames()
static ai.djl.ndarray.types.DataType	getDataType(com.sun.jna.Pointer ndArray)
static ai.djl.Device	getDevice(com.sun.jna.Pointer ndArray)
static java.util.Set<java.lang.String>	getFeatures()
static int	getGpuCount()
static long[]	getGpuMemory(ai.djl.Device device)
static com.sun.jna.Pointer	getGradient(com.sun.jna.Pointer handle)
static java.util.Map<java.lang.String,FunctionInfo>	getNdArrayFunctions()
static ai.djl.ndarray.types.Shape	getShape(com.sun.jna.Pointer ndArray)
static ai.djl.ndarray.types.SparseFormat	getStorageType(com.sun.jna.Pointer ndArray)
static com.sun.jna.Pointer	getSymbolInternals(com.sun.jna.Pointer symbol)
static com.sun.jna.Pointer	getSymbolOutput(com.sun.jna.Pointer symbol, int index)
static int	getVersion()
static ai.djl.util.PairList<com.sun.jna.Pointer,ai.djl.ndarray.types.SparseFormat>	imperativeInvoke(com.sun.jna.Pointer function, PointerArray inputs, com.sun.jna.ptr.PointerByReference destRef, ai.djl.util.PairList<java.lang.String,?> params)
static java.util.List<java.util.List<ai.djl.ndarray.types.Shape>>	inferShape(Symbol symbol, ai.djl.util.PairList<java.lang.String,ai.djl.ndarray.types.Shape> args)
static int	isNumpyMode()
static java.lang.String[]	listSymbolArguments(com.sun.jna.Pointer symbol)
static java.lang.String[]	listSymbolAuxiliaryStates(com.sun.jna.Pointer symbol)
static java.lang.String[]	listSymbolNames(com.sun.jna.Pointer symbol)
static java.lang.String[]	listSymbolOutputs(com.sun.jna.Pointer symbol)
static ai.djl.ndarray.NDList	loadNdArray(MxNDManager manager, java.nio.file.Path path, ai.djl.Device device)
static void	ndArraySyncCopyFromNdArray(MxNDArray dest, MxNDArray src, int location)
static FunctionInfo	op(java.lang.String opName)
static void	parameterStoreClose(com.sun.jna.Pointer handle)
static com.sun.jna.Pointer	parameterStoreCreate(java.lang.String type)
static void	parameterStoreInit(com.sun.jna.Pointer handle, int num, java.lang.String[] keys, ai.djl.ndarray.NDList vals)
static void	parameterStorePull(com.sun.jna.Pointer handle, int num, int[] keys, ai.djl.ndarray.NDList vals, int priority)
static void	parameterStorePull(com.sun.jna.Pointer handle, int num, java.lang.String[] keys, ai.djl.ndarray.NDList vals, int priority)
static void	parameterStorePush(com.sun.jna.Pointer handle, int num, java.lang.String[] keys, ai.djl.ndarray.NDList vals, int priority)
static void	parameterStoreSetUpdater(com.sun.jna.Pointer handle, MxnetLibrary.MXKVStoreUpdater updater, MxnetLibrary.MXKVStoreStrUpdater stringUpdater, com.sun.jna.Pointer updaterHandle)
static void	parameterStoreSetUpdater(com.sun.jna.Pointer handle, MxnetLibrary.MXKVStoreUpdater updater, com.sun.jna.Pointer updaterHandle)
static int	randomSeed(int seed)
static void	setNumpyMode(JnaUtils.NumpyMode mode)
static void	syncCopyFromCPU(com.sun.jna.Pointer ndArray, java.nio.Buffer data, int len)
static void	syncCopyToCPU(com.sun.jna.Pointer ndArray, com.sun.jna.Pointer data, int len)
static boolean	useThreadSafePredictor()
static void	waitAll()
static void	waitToRead(com.sun.jna.Pointer ndArray)
static void	waitToWrite(com.sun.jna.Pointer ndArray)

Methods inherited from class java.lang.Object

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

Field Detail

EMPTY_ARRAY

public static final java.lang.String[] EMPTY_ARRAY

MXNET_THREAD_SAFE_PREDICTOR

public static final java.lang.String MXNET_THREAD_SAFE_PREDICTOR

See Also:

Constant Field Values

Method Detail

getVersion

public static int getVersion()

getAllOpNames

public static java.util.Set<java.lang.String> getAllOpNames()

getNdArrayFunctions

public static java.util.Map<java.lang.String,FunctionInfo> getNdArrayFunctions()

op

public static FunctionInfo op(java.lang.String opName)

getGpuCount

public static int getGpuCount()

getGpuMemory

public static long[] getGpuMemory(ai.djl.Device device)

getFeatures

public static java.util.Set<java.lang.String> getFeatures()

randomSeed

public static int randomSeed(int seed)

createNdArray

public static com.sun.jna.Pointer createNdArray(ai.djl.Device device,
                                                ai.djl.ndarray.types.Shape shape,
                                                ai.djl.ndarray.types.DataType dtype,
                                                int size,
                                                boolean delayedAlloc)

createSparseNdArray

public static com.sun.jna.Pointer createSparseNdArray(ai.djl.ndarray.types.SparseFormat fmt,
                                                      ai.djl.Device device,
                                                      ai.djl.ndarray.types.Shape shape,
                                                      ai.djl.ndarray.types.DataType dtype,
                                                      ai.djl.ndarray.types.DataType[] auxDTypes,
                                                      ai.djl.ndarray.types.Shape[] auxShapes,
                                                      boolean delayedAlloc)

ndArraySyncCopyFromNdArray

public static void ndArraySyncCopyFromNdArray(MxNDArray dest,
                                              MxNDArray src,
                                              int location)

loadNdArray

public static ai.djl.ndarray.NDList loadNdArray(MxNDManager manager,
                                                java.nio.file.Path path,
                                                ai.djl.Device device)

freeNdArray

public static void freeNdArray(com.sun.jna.Pointer ndArray)

waitToRead

public static void waitToRead(com.sun.jna.Pointer ndArray)

waitToWrite

public static void waitToWrite(com.sun.jna.Pointer ndArray)

waitAll

public static void waitAll()

syncCopyToCPU

public static void syncCopyToCPU(com.sun.jna.Pointer ndArray,
                                 com.sun.jna.Pointer data,
                                 int len)

syncCopyFromCPU

public static void syncCopyFromCPU(com.sun.jna.Pointer ndArray,
                                   java.nio.Buffer data,
                                   int len)

imperativeInvoke

public static ai.djl.util.PairList<com.sun.jna.Pointer,ai.djl.ndarray.types.SparseFormat> imperativeInvoke(com.sun.jna.Pointer function,
                                                                                                           PointerArray inputs,
                                                                                                           com.sun.jna.ptr.PointerByReference destRef,
                                                                                                           ai.djl.util.PairList<java.lang.String,?> params)

getStorageType

public static ai.djl.ndarray.types.SparseFormat getStorageType(com.sun.jna.Pointer ndArray)

getDevice

public static ai.djl.Device getDevice(com.sun.jna.Pointer ndArray)

getShape

public static ai.djl.ndarray.types.Shape getShape(com.sun.jna.Pointer ndArray)

getDataType

public static ai.djl.ndarray.types.DataType getDataType(com.sun.jna.Pointer ndArray)

autogradSetIsRecording

public static boolean autogradSetIsRecording(boolean isRecording)

autogradSetTraining

public static boolean autogradSetTraining(boolean isTraining)

autogradIsRecording

public static boolean autogradIsRecording()

autogradIsTraining

public static boolean autogradIsTraining()

autogradMarkVariables

public static void autogradMarkVariables(int numVar,
                                         com.sun.jna.Pointer varHandles,
                                         java.nio.IntBuffer reqsArray,
                                         com.sun.jna.Pointer gradHandles)

autogradBackward

public static void autogradBackward(ai.djl.ndarray.NDList array,
                                    int retainGraph)

autogradBackwardExecute

public static void autogradBackwardExecute(int numOutput,
                                           ai.djl.ndarray.NDList array,
                                           ai.djl.ndarray.NDArray outgrad,
                                           int numVariables,
                                           com.sun.jna.Pointer varHandles,
                                           int retainGraph,
                                           int createGraph,
                                           int isTrain,
                                           com.sun.jna.Pointer gradHandles,
                                           com.sun.jna.Pointer gradSparseFormat)

autogradGetSymbol

public static com.sun.jna.Pointer autogradGetSymbol(ai.djl.ndarray.NDArray array)

isNumpyMode

public static int isNumpyMode()

setNumpyMode

public static void setNumpyMode(JnaUtils.NumpyMode mode)

getGradient

public static com.sun.jna.Pointer getGradient(com.sun.jna.Pointer handle)

parameterStoreCreate

public static com.sun.jna.Pointer parameterStoreCreate(java.lang.String type)

parameterStoreClose

public static void parameterStoreClose(com.sun.jna.Pointer handle)

parameterStoreInit

public static void parameterStoreInit(com.sun.jna.Pointer handle,
                                      int num,
                                      java.lang.String[] keys,
                                      ai.djl.ndarray.NDList vals)

parameterStorePush

public static void parameterStorePush(com.sun.jna.Pointer handle,
                                      int num,
                                      java.lang.String[] keys,
                                      ai.djl.ndarray.NDList vals,
                                      int priority)

parameterStorePull

public static void parameterStorePull(com.sun.jna.Pointer handle,
                                      int num,
                                      int[] keys,
                                      ai.djl.ndarray.NDList vals,
                                      int priority)

parameterStorePull

public static void parameterStorePull(com.sun.jna.Pointer handle,
                                      int num,
                                      java.lang.String[] keys,
                                      ai.djl.ndarray.NDList vals,
                                      int priority)

parameterStoreSetUpdater

public static void parameterStoreSetUpdater(com.sun.jna.Pointer handle,
                                            MxnetLibrary.MXKVStoreUpdater updater,
                                            MxnetLibrary.MXKVStoreStrUpdater stringUpdater,
                                            com.sun.jna.Pointer updaterHandle)

parameterStoreSetUpdater

public static void parameterStoreSetUpdater(com.sun.jna.Pointer handle,
                                            MxnetLibrary.MXKVStoreUpdater updater,
                                            com.sun.jna.Pointer updaterHandle)

getSymbolOutput

public static com.sun.jna.Pointer getSymbolOutput(com.sun.jna.Pointer symbol,
                                                  int index)

listSymbolOutputs

public static java.lang.String[] listSymbolOutputs(com.sun.jna.Pointer symbol)

freeSymbol

public static void freeSymbol(com.sun.jna.Pointer symbol)

listSymbolNames

public static java.lang.String[] listSymbolNames(com.sun.jna.Pointer symbol)

listSymbolArguments

public static java.lang.String[] listSymbolArguments(com.sun.jna.Pointer symbol)

listSymbolAuxiliaryStates

public static java.lang.String[] listSymbolAuxiliaryStates(com.sun.jna.Pointer symbol)

getSymbolInternals

public static com.sun.jna.Pointer getSymbolInternals(com.sun.jna.Pointer symbol)

createSymbolFromFile

public static com.sun.jna.Pointer createSymbolFromFile(java.lang.String path)

inferShape

public static java.util.List<java.util.List<ai.djl.ndarray.types.Shape>> inferShape(Symbol symbol,
                                                                                    ai.djl.util.PairList<java.lang.String,ai.djl.ndarray.types.Shape> args)

createCachedOp

public static CachedOp createCachedOp(MxSymbolBlock block,
                                      MxNDManager manager)

Creates cached op flags.

data_indices : [0, 2, 4] Used to label input location, param_indices : [1, 3] Used to label param location

Parameters:

block - the MxSymbolBlock that loaded in the backend

manager - the NDManager used to create NDArray

Returns:

a CachedOp for inference

freeCachedOp

public static void freeCachedOp(com.sun.jna.Pointer handle)

cachedOpInvoke

public static MxNDArray[] cachedOpInvoke(MxNDManager manager,
                                         com.sun.jna.Pointer cachedOpHandle,
                                         MxNDArray[] inputs)

useThreadSafePredictor

public static boolean useThreadSafePredictor()

checkCall

public static void checkCall(int ret)