def _decode(self, sequence, masked_positions):
batch_size = sequence.shape[<NUM_LIT:0>]<EOL>num_masked_positions = masked_positions.shape[<NUM_LIT:1>]<EOL>ctx = masked_positions.context<EOL>dtype = masked_positions.dtype<EOL>batch_idx = mx.nd.arange(<NUM_LIT:0>, batch_size, repeat=num_masked_positions, dtype=dtype, ctx=ctx)<EOL>batch_idx = batch_idx.reshape((<NUM_LIT:1>, -<NUM_LIT:1>))<EOL>masked_positions = masked_positions.reshape((<NUM_LIT:1>, -<NUM_LIT:1>))<EOL>position_idx = mx.nd.Concat(batch_idx, masked_positions, dim=<NUM_LIT:0>)<EOL>encoded = mx.nd.gather_nd(sequence, position_idx)<EOL>encoded = encoded.reshape((batch_size, num_masked_positions, sequence.shape[-<NUM_LIT:1>]))<EOL>decoded = self.decoder(encoded)<EOL>return decoded<EOL>
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).