Method InferShape

operator/operator/lstm.cpp:7–43

namespace TEngine {

bool LSTM::InferShape(const std::vector<TShape>& ishape, std::vector<TShape>& oshape, int layout)
{
    // input tensors:
    // 0 --- input: [seq_length, batch_size, input_size]
    // 1 --- kernel [ (input_size+hidden_size),cell_state_size]
    // others: optional

    // output tensor: [output_len, batch_size,hidden_size]
    // std::cout<<"!!!!!!!\n";
    const TShape input_shape = ishape[0];
    int batch_size = input_shape.Shape(1);
    if(param_.mxnet_flag == 0)
    {
        batch_size = input_shape.Shape(0);
    }

    //
    std::vector<int> dims(3);
    if(param_.mxnet_flag == 0)
    {
        dims[0] = batch_size;
        dims[1] = input_shape.Shape(0);
        dims[2] = param_.hidden_size;
    }
    else
    {
        dims[0] = input_shape.Shape(0);
        dims[1] = batch_size;
        dims[2] = param_.hidden_size;
    }

    // std::cout<<dims[0]<<","<< dims[1]<<","<<dims[2]<<"\n";

    oshape[0].SetDim(dims);

    return true;
}

void LSTM::SetSchema(void)
{