Incorrect LSTM backward pass implementation (iterates over input features instead of time steps)

[amedoc]

Hi,
Thanks for the huge effort on the project. However, there is a minor mistake in the LSTM node.
The current LSTM implementation seems to handle the backward pass incorrectly. After generated the code for the reverse direction, it should iterate over the time sequence dimension ("s" variable in the implementation), not the input features ("di"). This leads to wrong outputs for bidirectional or reverse LSTMs.

[kraiskil]

Thanks. There is currently one LSTM bidirectional test that passes.

I made a second test https://github.com/kraiskil/onnx2c/blob/95-incorrect-lstm-backward-pass-implementation-iterates-over-input-features-instead-of-time-steps/test/local_ops/test_lstm_bidirectional_issue_95/gentest.py, and this test fails with a wrong output. I'm not 100% sure this is the same error you are seeing. And going back to how LSTM works will take a moment.

This is the compiled code out of that test. What is the diff that would make this implementation correct?

/*
 * Operand:           LSTM
 * Name in ONNX file: anonymous_LSTM_0
 */
FUNC_PREFIX void node_anonymous_LSTM_0( const float X[1][3][2], const float W[2][16][2], const float R[2][16][4], float Y_h[2][3][4], float Y_c[2][3][4] )
{
	/* LSTM 
	 * inputs: 
	 *   X = tensor_X
	 *   W = tensor_W
	 *   R = tensor_R
	 *   B = 
	 *   sequence_lens = 
	 *   initial_h = 
	 *   initial_c = 
	 *   P = 
	 * outputs: 
	 *   Y = tensor_
	 *   Y_h = tensor_Y_h
	 *   Y_c = tensor_node_anonymous_LSTM_0_recursive_2
	 * attributes:
	 *   activations: Sigmoid Tanh Tanh Sigmoid Tanh Tanh 
	 * clip: off
	 * layout: 0
	 * (rest TBD):
	 */
	int hs = 4;
	int ds = 2;
	int bs = 3;
	int iidx = 0;
	int oidx = hs;
	int fidx = 2*hs;
	int cidx = 3*hs;
	int Rb = 4*hs;
	int sequence_lenght = 1;
	/* Forget gate */
	float ft[3][4];
	/* Input gate */
	float it[3][4];
	/* Cell gate */
	float ct[3][4];
	/* Output gate */
	float ot[3][4];

	memset(Y_h, 0, sizeof(*Y_h));
	memset(Y_c, 0, sizeof(*Y_c));

	for( int s=0; s<sequence_lenght; s++) {

		/* Forward lane */
		for( int b=0; b<bs; b++)
		for( int h=0; h<hs; h++) {
			ft[b][h]=0;
			it[b][h]=0;
			ct[b][h]=0;
			for( int i=0; i<ds; i++) {
				ft[b][h] += X[s][b][i]*W[0][fidx+h][i];
				it[b][h] += X[s][b][i]*W[0][iidx+h][i];
				ct[b][h] += X[s][b][i]*W[0][cidx+h][i];
			}
			for( int k=0; k<hs; k++) {
				ft[b][h] += Y_h[0][b][k]*R[0][fidx+h][k];
				ct[b][h] += Y_h[0][b][k]*R[0][cidx+h][k];
				it[b][h] += Y_h[0][b][k]*R[0][iidx+h][k];
			}
			ft[b][h] =1.0f/(1+expf(-ft[b][h]));
			it[b][h] =1.0f/(1+expf(-it[b][h]));
			ct[b][h] =tanh(ct[b][h]);
		}
		for( int b=0; b<bs; b++)
		for( int h=0; h<hs; h++) {
			/* Cell state */
			Y_c[0][b][h] = Y_c[0][b][h]*ft[b][h] + it[b][h]*ct[b][h];
			/* Output gate */
			ot[b][h]=0;
			for( int i=0; i<ds; i++)
				ot[b][h] += X[s][b][i]*W[0][oidx+h][i];
			for( int k=0; k<hs; k++)
				ot[b][h] += Y_h[0][b][k]*R[0][oidx+h][k];
			ot[b][h] =1.0f/(1+expf(-ot[b][h]));
		}
		/* Hidden state */
		for( int b=0; b<bs; b++)
		for( int h=0; h<hs; h++) {
			Y_h[0][b][h] = ot[b][h] * tanh(Y_c[0][b][h]);
		}


		/* Backward lane */
		for( int b=0; b<bs; b++)
		for( int h=0; h<hs; h++) {
			ft[b][h]=0;
			it[b][h]=0;
			ct[b][h]=0;
			for( int i=0; i<ds; i++) {
				ft[b][h] += X[s][b][ds-1-i]*W[1][fidx+h][i];
				it[b][h] += X[s][b][ds-1-i]*W[1][iidx+h][i];
				ct[b][h] += X[s][b][ds-1-i]*W[1][cidx+h][i];
			}
			for( int k=0; k<hs; k++) {
				ft[b][h] += Y_h[1][b][k]*R[1][fidx+h][k];
				ct[b][h] += Y_h[1][b][k]*R[1][cidx+h][k];
				it[b][h] += Y_h[1][b][k]*R[1][iidx+h][k];
			}
			ft[b][h] =1.0f/(1+expf(-ft[b][h]));
			it[b][h] =1.0f/(1+expf(-it[b][h]));
			ct[b][h] =tanh(ct[b][h]);
		}
		for( int b=0; b<bs; b++)
		for( int h=0; h<hs; h++) {
			/* Cell state */
			Y_c[1][b][h] = Y_c[1][b][h]*ft[b][h] + it[b][h]*ct[b][h];
			/* Output gate */
			ot[b][h]=0;
			for( int i=0; i<ds; i++)
				ot[b][h] += X[s][b][ds-1-i]*W[1][oidx+h][i];
			for( int k=0; k<hs; k++)
				ot[b][h] += Y_h[1][b][k]*R[1][oidx+h][k];
			ot[b][h] =1.0f/(1+expf(-ot[b][h]));
		}
		/* Hidden state */
		for( int b=0; b<bs; b++)
		for( int h=0; h<hs; h++) {
			Y_h[1][b][h] = ot[b][h] * tanh(Y_c[1][b][h]);
		}

	} /* sequences */
}