%% Create Bidirectional LSTM (BiLSTM) Function
% This example shows how to create a bidirectional long-short term memory (BiLSTM) function for dlarray objects.
% In a deep learning model, a bidirectional LSTM (BiLSTM) operation learns bidirectional long-term dependencies between time steps of time series or sequence data. These dependencies can be useful when you want the network to learn from the complete time series at each time step.
% For most tasks, you can train a network using the trainNetwork function with using a bilstmLayer object. To train a BiLSTM network with a custom training loop, you can create a custom BiLSTM function using this example as a guide.
% A BiLSTM consists of two LSTM components: the "forward LSTM" that operates from the first time step to the last time step and the "backward LSTM" that operates from the last time step to the first time step. After passing the data through the two LSTM components, the operation concatenates the outputs together along the channel dimension.
%
%% Create BiLSTM Function
% Create the bilstm function, listed at the end of the example, that applies a BiLSTM operation to the input dlX using the initial hidden state H0, initial cell state C0, and parameters weights, recurrentWeights, and bias.
%% Initialize BiLSTM Parameters
% Specify the input size (for example, the embedding dimension of the input layer) and the number of hidden units.

inputSize = 256;
numHiddenUnits = 50;

% Initialize the BILSTM parameters. The BiLSTM operation requires a set of input weights, recurrent weights, and bias for both the forward and backward parts of the operation. For these parameters, specify the concatenation of the forward and backward components. In this case, the input weights have size [8*numHiddenUnits inputSize], the recurrent weights have size [8*numHiddenUnits numHiddenUnits], and the bias has size [8*numHiddenUnits 1].

inputWeights = dlarray(0.01*randn([8*numHiddenUnits inputSize],'single'));
recurrentWeights = dlarray(0.01*randn([8*numHiddenUnits numHiddenUnits],'single'));
bias = dlarray(0.01*randn([8*numHiddenUnits 1],'single'));

% Initialize the BiLSTM hidden and cell state. Similar to the parameters, specify the concatenation of the forward and backward components. In this case, the hidden and cell state each have size [2*numHiddenUnits 1].

H0 = dlarray(zeros(2*numHiddenUnits,1,'single'));
C0 = dlarray(zeros(2*numHiddenUnits,1,'single'));

%% Apply BiLSTM Operation
% Specify an array of random data with mini-batch size 128 and sequence length 75. The first dimension of the input (the channel dimension) must match the input size of the BiLSTM operation.

miniBatchSize = 128;
sequenceLength = 75;
dlX = dlarray(randn([inputSize miniBatchSize sequenceLength],'single'),'CBT');

% Apply the BiLSTM operation and view the size of the output.

dlY = bilstm(dlX,H0,C0,inputWeights,recurrentWeights,bias);
size(dlY)

% For models that require only the last time step of the sequence, extract the vectors corresponding to the last output of the forward LSTM and backward LSTM components.

dlYLastForward = dlY(1:numHiddenUnits,:,end);
dlYLastBackward = dlY(numHiddenUnits+1:end,:,1);

dlYLast = cat(1, dlYLastForward, dlYLastBackward);
size(dlYLast)

%
%% BiLSTM Function
% The bilstm function applies a BiLSTM operation to the input dlX using the initial hidden state H0, initial cell state C0, and parameters weights, recurrentWeights, and bias. The input weights have size [8*numHiddenUnits inputSize], the recurrent weights have size [8*numHiddenUnits numHiddenUnits], and the bias has size [8*numHiddenUnits 1]. The hidden and cell state each have size [2*numHiddenUnits 1].

function dlY = bilstm(dlX,H0,C0,inputWeights,recurrentWeights,bias)

% Determine forward and backward parameter indices
numHiddenUnits = numel(bias)/8;
idxForward = 1:4*numHiddenUnits;
idxBackward = 4*numHiddenUnits+1:8*numHiddenUnits;

% Forward and backward states
H0Forward = H0(1:numHiddenUnits);
H0Backward = H0(numHiddenUnits+1:end);
C0Forward = C0(1:numHiddenUnits);
C0Backward = C0(numHiddenUnits+1:end);

% Forward and backward parameters
inputWeightsForward = inputWeights(idxForward,:);
inputWeightsBackward = inputWeights(idxBackward,:);
recurrentWeightsForward = recurrentWeights(idxForward,:);
recurrentWeightsBackward = recurrentWeights(idxBackward,:);
biasForward = bias(idxForward);
biasBackward = bias(idxBackward);

% Forward LSTM
dlYForward = lstm(dlX,H0Forward,C0Forward,inputWeightsForward, ...
    recurrentWeightsForward,biasForward);

% Backward LSTM
dlXBackward = flip(dlX,3);
dlYBackward = lstm(dlXBackward,H0Backward,C0Backward,inputWeightsBackward, ...
    recurrentWeightsBackward,biasBackward);
dlYBackward = flip(dlYBackward,3);

% Output
dlY = cat(1,dlYForward,dlYBackward);

end

% Copyright 2019 The MathWorks, Inc.