ai.djl.nn

Class LambdaBlock

java.lang.Object

ai.djl.nn.AbstractBlock

ai.djl.nn.LambdaBlock

All Implemented Interfaces:

Block

public class LambdaBlock
extends AbstractBlock

LambdaBlock is a Block with no parameters or children.

LambdaBlock allows converting any function that takes an NDList as input and returns an NDList into a parameter-less and child-less Block. This can be useful in converting activation functions, identity blocks, and more. For example, identity block can be created as new LambdaBlock(x -> x).

Field Summary

Fields inherited from class ai.djl.nn.AbstractBlock

children, inputNames, inputShapes, parameters, version

Constructor Summary

Constructors

Constructor and Description
LambdaBlock(java.util.function.Function<NDList,NDList> lambda) Creates a LambdaBlock that can apply the specified function.

Method Summary

Modifier and Type	Method and Description
protected NDList	forwardInternal(ParameterStore parameterStore, NDList inputs, boolean training, ai.djl.util.PairList<java.lang.String,java.lang.Object> params) A helper for Block.forward(ParameterStore, NDList, boolean, PairList) after initialization.
Shape[]	getOutputShapes(Shape[] inputShapes) Returns the expected output shapes of the block for the specified input shapes.
void	loadParameters(NDManager manager, java.io.DataInputStream is) Loads the parameters from the given input stream.
static LambdaBlock	singleton(java.util.function.Function<NDArray,NDArray> lambda) Creates a LambdaBlock for a singleton function.
java.lang.String	toString()

Methods inherited from class ai.djl.nn.AbstractBlock

addChildBlock, addParameter, beforeInitialize, cast, clear, describeInput, forward, forward, forwardInternal, getChildren, getDirectParameters, getParameters, initialize, initializeChildBlocks, isInitialized, loadMetadata, prepare, readInputShapes, saveInputShapes, saveMetadata, saveParameters, setInitializer, setInitializer, setInitializer

Methods inherited from class java.lang.Object

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

Methods inherited from interface ai.djl.nn.Block

forward, validateLayout

Constructor Detail

LambdaBlock

public LambdaBlock(java.util.function.Function<NDList,NDList> lambda)

Creates a LambdaBlock that can apply the specified function.

Parameters:

lambda - the function to apply

Method Detail

singleton

public static LambdaBlock singleton(java.util.function.Function<NDArray,NDArray> lambda)

Creates a LambdaBlock for a singleton function.

Parameters:

lambda - a function accepting a singleton NDList and returning another sinleton NDList

Returns:

a new LambdaBlock for the function

forwardInternal

protected NDList forwardInternal(ParameterStore parameterStore,
                                 NDList inputs,
                                 boolean training,
                                 ai.djl.util.PairList<java.lang.String,java.lang.Object> params)

A helper for Block.forward(ParameterStore, NDList, boolean, PairList) after initialization.

Specified by:

forwardInternal in class AbstractBlock

Parameters:

parameterStore - the parameter store

inputs - the input NDList

training - true for a training forward pass

params - optional parameters

Returns:

the output of the forward pass

getOutputShapes

public Shape[] getOutputShapes(Shape[] inputShapes)

Returns the expected output shapes of the block for the specified input shapes.

Parameters:

inputShapes - the shapes of the inputs

Returns:

the expected output shapes of the block

loadParameters

public void loadParameters(NDManager manager,
                           java.io.DataInputStream is)
                    throws java.io.IOException,
                           MalformedModelException

Loads the parameters from the given input stream.

Specified by:

loadParameters in interface Block

Overrides:

loadParameters in class AbstractBlock

Parameters:

manager - an NDManager to create the parameter arrays

is - the inputstream that stream the parameter values

Throws:

java.io.IOException - if an I/O error occurs

MalformedModelException - if the model file is corrupted or unsupported

toString

public java.lang.String toString()

Overrides:

toString in class AbstractBlock