// ApplyBlock applies the LSTM to a batch of inputs.
func (l *LSTM) ApplyBlock(s []State, in []autofunc.Result) BlockResult {
var internalPool, lastOutPool []*autofunc.Variable
res := autofunc.PoolAll(in, func(in []autofunc.Result) autofunc.Result {
var weavedInputs []autofunc.Result
var internalResults []autofunc.Result
for i, sObj := range s {
state := sObj.(lstmState)
internalVar := &autofunc.Variable{Vector: state.Internal}
lastOutVar := &autofunc.Variable{Vector: state.Output}
internalPool = append(internalPool, internalVar)
lastOutPool = append(lastOutPool, lastOutVar)
weavedInputs = append(weavedInputs, in[i], lastOutVar, internalVar)
internalResults = append(internalResults, internalVar)
}
gateIn := autofunc.Concat(weavedInputs...)
inValue := l.inputValue.Batch(gateIn, len(in))
inGate := l.inputGate.Batch(gateIn, len(in))
rememberGate := l.rememberGate.Batch(gateIn, len(in))
lastState := autofunc.Concat(internalResults...)
newState := autofunc.Add(autofunc.Mul(rememberGate, lastState),
autofunc.Mul(inValue, inGate))
return autofunc.Pool(newState, func(newState autofunc.Result) autofunc.Result {
var newWeaved []autofunc.Result
for i, state := range autofunc.Split(len(in), newState) {
newWeaved = append(newWeaved, in[i], lastOutPool[i], state)
}
newGateIn := autofunc.Concat(newWeaved...)
outGate := l.outputGate.Batch(newGateIn, len(in))
outValues := neuralnet.HyperbolicTangent{}.Apply(newState)
return autofunc.Concat(newState, autofunc.Mul(outGate, outValues))
})
})
states, outs := splitLSTMOutput(len(in), res.Output())
return &lstmResult{
CellStates: states,
OutputVecs: outs,
InternalPool: internalPool,
LastOutPool: lastOutPool,
JoinedOut: res,
}
}
func (l *lstmGate) Batch(in autofunc.Result, n int) autofunc.Result {
if l.Peephole == nil {
return l.Activation.Apply(l.Dense.Batch(in, n))
}
return autofunc.Pool(in, func(in autofunc.Result) autofunc.Result {
vecSize := len(in.Output()) / n
var weightedInputs []autofunc.Result
var peepholed []autofunc.Result
for i := 0; i < n; i++ {
start := vecSize * i
weightedEnd := start + vecSize - len(l.Peephole.Vector)
weightedInputs = append(weightedInputs, autofunc.Slice(in, start, weightedEnd))
peepholeMe := autofunc.Slice(in, weightedEnd, (i+1)*vecSize)
peepholed = append(peepholed, autofunc.Mul(l.Peephole, peepholeMe))
}
weighted := l.Dense.Batch(autofunc.Concat(weightedInputs...), n)
return l.Activation.Apply(autofunc.Add(autofunc.Concat(peepholed...), weighted))
})
}
// ApplyBlock applies the block to an input.
func (g *GRU) ApplyBlock(s []State, in []autofunc.Result) BlockResult {
stateVars, stateRes := PoolVecStates(s)
var gateInputs []autofunc.Result
for i, x := range stateRes {
gateInputs = append(gateInputs, in[i], x)
}
n := len(in)
gateInput := autofunc.Concat(gateInputs...)
stateIn := autofunc.Concat(stateRes...)
resetMask := g.resetGate.Batch(gateInput, n)
updateMask := g.updateGate.Batch(gateInput, n)
maskedByReset := autofunc.Mul(resetMask, stateIn)
inputValue := autofunc.PoolSplit(n, maskedByReset,
func(newStates []autofunc.Result) autofunc.Result {
var newGateInputs []autofunc.Result
for i, input := range in {
newGateInputs = append(newGateInputs, input, newStates[i])
}
newIn := autofunc.Concat(newGateInputs...)
return g.inputValue.Batch(newIn, n)
})
newState := autofunc.Pool(updateMask, func(umask autofunc.Result) autofunc.Result {
updateComplement := autofunc.AddScaler(autofunc.Scale(umask, -1), 1)
return autofunc.Add(autofunc.Mul(umask, stateIn),
autofunc.Mul(updateComplement, inputValue))
})
return &gruResult{
InStates: stateVars,
Output: newState,
}
}
