Package io.github.kirstenali.deepj.layers.transformer.attention

Class MultiHeadSelfAttention

java.lang.Object

io.github.kirstenali.deepj.layers.transformer.attention.MultiHeadSelfAttention

All Implemented Interfaces:

Layer, Trainable

Direct Known Subclasses:

RoPEMultiHeadSelfAttention

public class MultiHeadSelfAttention
extends Object
implements Layer

Field Summary

Fields

Modifier and Type	Field	Description
protected final int	headDim
protected final int	nHeads

Constructor Summary

Constructors

Constructor	Description
MultiHeadSelfAttention(int dModel, int nHeads, boolean causalMask, Random rnd)
MultiHeadSelfAttention(int dModel, int nHeads, boolean causalMask, Random rnd, ParameterFactory parameters)

Method Summary

Modifier and Type	Method	Description
Tensor	forward(Tensor x)
Tensor	forward(Tensor input, InferenceContext context)
List<Parameter>	parameters()
boolean	supportsIncrementalInference()
protected Tensor	transformQueryKey(Tensor heads, int seqLen)
protected Tensor	transformQueryKey(Tensor heads, int length, int offset)
protected Tensor	transformQueryKey(Tensor heads, int length, InferenceContext context)

Methods inherited from class java.lang.Object

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

Methods inherited from interface io.github.kirstenali.deepj.training.Trainable

zeroGrad

Field Details

nHeads

protected final int nHeads

headDim

protected final int headDim

Constructor Details

MultiHeadSelfAttention

public MultiHeadSelfAttention(int dModel,
 int nHeads,
 boolean causalMask,
 Random rnd)

MultiHeadSelfAttention

public MultiHeadSelfAttention(int dModel,
 int nHeads,
 boolean causalMask,
 Random rnd,
 ParameterFactory parameters)

Method Details

forward

public Tensor forward(Tensor x)

Specified by:

forward in interface Layer

transformQueryKey

protected Tensor transformQueryKey(Tensor heads,
 int seqLen)

transformQueryKey

protected Tensor transformQueryKey(Tensor heads,
 int length,
 int offset)

transformQueryKey

protected Tensor transformQueryKey(Tensor heads,
 int length,
 InferenceContext context)

supportsIncrementalInference

public boolean supportsIncrementalInference()

Specified by:

supportsIncrementalInference in interface Layer

forward

public Tensor forward(Tensor input,
 InferenceContext context)

Specified by:

forward in interface Layer

parameters

public List<Parameter> parameters()

Specified by:

parameters in interface Trainable