How to set Cell State for two IRNNv2Layer when input feature size is not equal to hidden size?

AI & Data Science Deep Learning (Training & Inference) TensorRT
1

Hello everyone,

/*************************************************************************

Platform: TensorRT 5.1

*************************************************************************/

I am making a RNN network with TensorRT.

It has two LSTM layers.

The input is in Dim(1, 13, 512). And the hidden size is 256.

so dimension of each is:
c(0): Dim( 512 )
h(0): Dim( 256 )
c(1): Dim( 256 )
h(1): Dim( 256 ), where 0 is layer0, 1 is layer1

#define RNN_NAME_HIDDEN_INPUT "rnn/hidden_in"
#define RNN_NAME_CELL_INPUT   "rnn/cell_in"
#define RNN_LAYER_COUNT       2
#define RNN_HIDDEN_NUM        256
#define RNN_SEQUEN_LEN        13
IRNNv2Layer* addRNNv2Layer(ITensor* features, INetworkDefinition* network, std::map<std::string, Weights>& WeightsMap)
{
    ITensor* hiddenIn = network->addInput(RNN_NAME_HIDDEN_INPUT, DataType::kFLOAT, Dims3(1, RNN_LAYER_COUNT, RNN_HIDDEN_NUM));
    assert(hiddenIn != nullptr);

    ITensor* cellIn = network->addInput(RNN_NAME_CELL_INPUT, DataType::kFLOAT, Dims3(1, RNN_LAYER_COUNT, RNN_HIDDEN_NUM));
    assert(cellIn != nullptr);

    Dims dim;
    dim.nbDims = 1;
    dim.d[0] = 1;
    ITensor* seqLenIn = network->addInput(RNN_NAME_SEQ_LEN, DataType::kINT32, dim);
    assert(seqLenIn != nullptr);

    IRNNv2Layer* rnn = network->addRNNv2(*features, RNN_LAYER_COUNT, RNN_HIDDEN_NUM, RNN_SEQUEN_LEN, RNNOperation::kLSTM);
    assert(rnn != nullptr);

    ITensor* rnnOutput = rnn->getOutput(0);
    rnnOutput->setName("RNN_Output");

    rnn->setHiddenState(*hiddenIn);
    if(rnn->getOperation() == RNNOperation::kLSTM)
    {
        rnn->setCellState(*cellIn);
    }

    rnn->setSequenceLengths(*seqLenIn);

    seqLenIn->setLocation(TensorLocation::kDEVICE);
    ...
}

When I use rnn->setCellState(*cellIn) to specify the initial state of Cell. It seems the shape of the cell in two layer must be the same.

Is there any way to specify a different cell state for the RNN Layers?