def forward(self, prev_samples, upper_tier_conditioning, state=None):
        """
        prev_samples: (B, M_i, D_MODEL) if self.reproduce else (B, M_i, FRAME, D_MODEL)
        upper_tier_conditioning: (B, M_i, D_HIDDEN) or None
        """
        if not self.reproduce:
            # Use strided convolutions to get frame embeddings
            # This generalizes the SampleRNN operation to handle non-1D signals
            # This reshapes from (B, M_i, FRAME, D_MODEL) -> (B, M_i, D_HIDDEN)
            prev_samples = prev_samples.view(prev_samples.shape[0], -1, self.d_model)
            input = self.input_expand(prev_samples.permute(0, 2, 1)).permute(0, 2, 1)
        else:
            # SampleRNN uses an MLP (implemented as 1D Conv) to map (FRAME_SIZE, 1) to D_HIDDEN
            # This reshapes from (B, M_i, FRAME) -> (B, M_i, D_HIDDEN)
            input = self.input_expand(prev_samples.permute(0, 2, 1)).permute(0, 2, 1)

        if upper_tier_conditioning is not None:
            input += upper_tier_conditioning

        # Run RNN: (B, M_i, D_HIDDEN) -> (B, M_i, D_HIDDEN)
        if self.layer == 'gru':
            output, state = self.rnn(input, state.contiguous())
        elif self.layer == 's4':
            # TODO: not working
            output, state = self.rnn(input.transpose(1, 2), state)
            output = output.transpose(1, 2)

        # Run 1D transposed convolution to upsample: (B, M_i, D_HIDDEN) -> (B, M', D_HIDDEN)
        # TODO: make M' more precise
        output = self.upsampling(output.permute(0, 2, 1)).permute(0, 2, 1)

        return output, state

        prev_samples: (B, M_i, D_MODEL) if self.reproduce else (B, M_i, FRAME, D_MODEL)
        upper_tier_conditioning: (B, M_i, D_HIDDEN) or None
        
analyze/convnexts4nd/src/models/baselines/samplernn.py