docs: def forward(self, X):
          """Applies one step of an lstm (plus reduction) to the input X, which
          is handled by self.encoder"""
          # Identify the first non-zero integer from the right (i.e., the length
          # of the sequence before padding starts).
          batch_size, max_seq = X.shape[0], X.shape[1]
          seq_lengths = torch.zeros(batch_size, dtype=torch.long)
          for i in range(batch_size):
              for j in range(max_seq - 1, -1, -1):
                  if not torch.all(X[i, j] == 0):
                      seq_lengths[i] = j + 1
                      break
  
          # Sort by length because pack_padded_sequence requires it
          # Save original order to restore before returning
          seq_lengths, perm_idx = seq_lengths.sort(0, descending=True)
          X = X[perm_idx]
          inv_perm_idx = torch.tensor(
              [i for i, _ in sorted(enumerate(perm_idx), key=lambda idx: idx[1])],
              dtype=torch.long,
          )
  
          # Encode and pack input sequence
          X_packed = rnn_utils.pack_padded_sequence(
              self.encoder.encode(X), seq_lengths, batch_first=True
          )
  
          # Run LSTM
          outputs, (h_t, c_t) = self.lstm(X_packed)
  
          # Unpack and reduce outputs
          outputs_unpacked, _ = rnn_utils.pad_packed_sequence(outputs, batch_first=True)
          reduced = self._reduce_output(outputs_unpacked, seq_lengths)
          return reduced[inv_perm_idx, :]