def forward(self, inputs, *args, state=None, **kwargs):
        """
        inputs shape: (B, L, D) e.g. (16, 16447, 1)

        For SampleRNN, inputs contains quantized audio samples (e.g. B elements of length L)
        """
        batch_size = inputs.shape[0]

        assert state is not None, "SampleRNN: State should be constructed with default_state before forward pass"

        upper_tier_conditioning = None
        new_states = []
        for rnn, state_ in zip(reversed(self.frame_level_rnns), reversed(state)):
            # TODO: explain this
            from_index = self.lookback - rnn.n_frame_samples
            to_index = -rnn.n_frame_samples + 1

            # prev_samples shape: (B, CHUNK_SIZE, D) e.g. (16, 16384, 1)
            prev_samples = inputs[:, from_index : to_index]

            if self.reproduce:
                # SampleRNN dequantizes to recover the raw audio signal before passing this to the RNN
                prev_samples = self.dequantizer(prev_samples, bits=self.bits)
                prev_samples = 2 * prev_samples.contiguous()
                # Below, reshape from (B, CHUNK_SIZE, D) -> (B, -1, rnn.n_frame_samples) = (B, M_i, F_i)
                # e.g. (16, 16384, 1) -> (16, 256, 64) [first rnn] | (16, 1024, 16) [second rnn]
                prev_samples = prev_samples.view(batch_size, -1, rnn.n_frame_samples)

            else:
                # More generally, we can use an Embedding encoder instead
                prev_samples = self.encoder(prev_samples)
                prev_samples = prev_samples.contiguous()
                prev_samples = prev_samples.view(batch_size, -1, rnn.n_frame_samples, self.d_model)

            # upper_tier_conditioning shape: None -> (B, M, D_HIDDEN) [first rnn]
            # (B, M_{i-1}, D_HIDDEN) -> (B, M_i, D_HIDDEN) [second rnn]
            upper_tier_conditioning, new_state = rnn(prev_samples, upper_tier_conditioning, state_)

            new_states.append(new_state)

        # Make sure new states are in the right order
        new_states = list(reversed(new_states))

        bottom_frame_size = self.frame_level_rnns[0].frame_size
        mlp_input_sequences = inputs[:, self.lookback - bottom_frame_size : ]

        # mlp_input_sequences shape: (B, L - _, D) e.g. (16, 16399, 1)
        # upper_tier_conditioning shape: (B, M_{last_rnn}, D_HIDDEN) [last rnn]
        return self.sample_level_mlp(mlp_input_sequences, upper_tier_conditioning), new_states

        inputs shape: (B, L, D) e.g. (16, 16447, 1)

        For SampleRNN, inputs contains quantized audio samples (e.g. B elements of length L)
        
analyze/convnexts4nd/src/models/baselines/samplernn.py