ai.djl.nn.transformer

Class MemoryScope

java.lang.Object

ai.djl.nn.transformer.MemoryScope

All Implemented Interfaces:

java.lang.AutoCloseable

public final class MemoryScope
extends java.lang.Object
implements java.lang.AutoCloseable

Helper class for more complicated memory management scenarios. Allows to avoid boilerplate for memory handling. Makes sure the sub NDManager used is connected to the correct GPU to avoid crashes.

Method Summary

Modifier and Type	Method and Description
MemoryScope	add(NDArray... arrays) Adds the given arrays to this scopes sub manager.
MemoryScope	add(NDList... lists) Adds all arrays in the given lists to this memory scope.
void	close() Closes this scope by closing the sub manager used to manage it.
static MemoryScope	from(NDArray ndArray) Creates a new memory scope for the device of the given array and adds the array.
static MemoryScope	from(NDList list) Creates a new memory scope that fits the device of the first array in the given list, adds all arrays in the given list.
NDManager	getScopeManager() Returns the NDManager used to manage this scopes resources.
MemoryScope	remove(NDArray... arrays) Remove the given arrays from this scope and attach them back to this scopes parent NDManager.
MemoryScope	remove(NDList... lists) Remove the given arrays from this scope and attach them back to this scopes parent NDManager.
MemoryScope	waitToRead(NDArray... arrays) Waits for all given arrays to be ready to read, i.e.
MemoryScope	waitToRead(NDList... lists) Waits for all arrays in all given lists to be ready to be read, i.e.

Methods inherited from class java.lang.Object

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

Method Detail

add

public MemoryScope add(NDList... lists)

Adds all arrays in the given lists to this memory scope.

Parameters:

lists - the lists whose arrays to add to this scope, may be empty

Returns:

this scope

add

public MemoryScope add(NDArray... arrays)

Adds the given arrays to this scopes sub manager.

Parameters:

arrays - the arrays to add

Returns:

this scope

remove

public MemoryScope remove(NDList... lists)

Remove the given arrays from this scope and attach them back to this scopes parent NDManager.

Parameters:

lists - the lists containing the arrays to remove

Returns:

this scope

remove

public MemoryScope remove(NDArray... arrays)

Remove the given arrays from this scope and attach them back to this scopes parent NDManager.

Parameters:

arrays - arrays to remove

Returns:

this scope

getScopeManager

public NDManager getScopeManager()

Returns the NDManager used to manage this scopes resources.

Returns:

the NDManager used to manage this scopes resources

waitToRead

public MemoryScope waitToRead(NDArray... arrays)

Waits for all given arrays to be ready to read, i.e. waits for pending computations that write to them, then removes them from this scope.

Parameters:

arrays - arrays to wait for

Returns:

this scope

waitToRead

public MemoryScope waitToRead(NDList... lists)

Waits for all arrays in all given lists to be ready to be read, i.e. waits for pending computations that write to them, then removes them from this scope.

Parameters:

lists - may be empty

Returns:

this scope

close

public void close()

Closes this scope by closing the sub manager used to manage it. This causes all arrays still attached to this scope to be closed as well.

Specified by:

close in interface java.lang.AutoCloseable

from

public static MemoryScope from(NDArray ndArray)

Creates a new memory scope for the device of the given array and adds the array.

Parameters:

ndArray - an array

Returns:

a new memory scrope containing the array

from

public static MemoryScope from(NDList list)

Creates a new memory scope that fits the device of the first array in the given list, adds all arrays in the given list.

Parameters:

list - a list of arrays, must not be empty

Returns:

a new memory scope