I'm trying to use the provided LSTM layer implementation, but it seems to be incorrect.

[pmingkr]

djl/api/src/main/java/ai/djl/nn/recurrent/LSTM.java

Line 68 in 2740c21

if (inputs.size() == 1) {

// LSTM internal implementation

    /** {@inheritDoc} */
    @Override
    protected NDList forwardInternal(
            ParameterStore parameterStore,
            NDList inputs,
            boolean training,
            PairList<String, Object> params) {
        NDArrayEx ex = inputs.head().getNDArrayInternal();
        Device device = inputs.head().getDevice();
        NDList rnnParams = new NDList();
        for (Parameter parameter : parameters.values()) {
            rnnParams.add(parameterStore.getValue(parameter, device, training));
        }

        NDArray input = inputs.head();
        if (inputs.size() == 1) {
            int batchIndex = batchFirst ? 0 : 1;
            Shape stateShape =
                    new Shape(
                            (long) numLayers * getNumDirections(),
                            input.size(batchIndex),
                            stateSize);
            // hidden state
            inputs.add(input.getManager().zeros(stateShape));
            // cell
            inputs.add(input.getManager().zeros(stateShape));
        }
        if (inputs.size() == 2) {
            int batchIndex = batchFirst ? 0 : 1;
            Shape stateShape =
                    new Shape(
                            (long) numLayers * getNumDirections(),
                            input.size(batchIndex),
                            stateSize);
            // cell
            inputs.add(input.getManager().zeros(stateShape));
        }
        NDList outputs =
                ex.lstm(
                        input,
                        new NDList(inputs.get(1), inputs.get(2)),
                        rnnParams,
                        hasBiases,
                        numLayers,
                        dropRate,
                        training,
                        bidirectional,
                        batchFirst);
        if (returnState) {
            return outputs;
        }
        outputs.stream().skip(1).forEach(NDArray::close);
        return new NDList(outputs.get(0));
    }

it seems it always makes zero cells if the multiple inputs are not provided.

djl/examples/src/main/java/ai/djl/examples/training/TrainMnistWithLSTM.java

Line 87 in 2740c21

block.addSingleton(

and the example always makes a singleton list on reshaping.
so, cells are always zero?
I think it appears to be incorrect. because it does not memorize anything.