code: def compute_lengths(sequence_data: mx.sym.Symbol) -> mx.sym.Symbol:
      
      return mx.sym.sum(sequence_data != C.PAD_ID, axis=1)
docs: Computes sequence lengths of PAD_ID-padded data in sequence_data.
  
      :param sequence_data: Input data. Shape: (batch_size, seq_len).
      :return: Length data. Shape: (batch_size,).