code: def _decode(self, sequence, masked_positions):
          
          batch_size = sequence.shape[0]
          num_masked_positions = masked_positions.shape[1]
          ctx = masked_positions.context
          dtype = masked_positions.dtype
          # batch_idx = [0,0,0,1,1,1,2,2,2...]
          # masked_positions = [1,2,4,0,3,4,2,3,5...]
          batch_idx = mx.nd.arange(0, batch_size, repeat=num_masked_positions, dtype=dtype, ctx=ctx)
          batch_idx = batch_idx.reshape((1, -1))
          masked_positions = masked_positions.reshape((1, -1))
          position_idx = mx.nd.Concat(batch_idx, masked_positions, dim=0)
          encoded = mx.nd.gather_nd(sequence, position_idx)
          encoded = encoded.reshape((batch_size, num_masked_positions, sequence.shape[-1]))
          decoded = self.decoder(encoded)
          return decoded
docs: Generate unnormalized prediction for the masked language model task.
  
          This is only used for pre-training the BERT model.
  
          Inputs:
              - **sequence**: input tensor of sequence encodings.
                Shape (batch_size, seq_length, units).
              - **masked_positions**: input tensor of position of tokens for masked LM decoding.
                Shape (batch_size, num_masked_positions). For each sample in the batch, the values
                in this tensor must not be out of bound considering the length of the sequence.
  
          Outputs:
              - **masked_lm_outputs**: output tensor of token predictions for target masked_positions.
                  Shape (batch_size, num_masked_positions, vocab_size).