def forward(self, inputs, states=None, valid_length=None, steps=None): 
length = inputs.shape[<NUM_LIT:1>]<EOL>if valid_length is not None:<EOL><INDENT>arange = mx.nd.arange(length, ctx=valid_length.context, dtype=valid_length.dtype)<EOL>mask = mx.nd.broadcast_lesser(<EOL>arange.reshape((<NUM_LIT:1>, -<NUM_LIT:1>)),<EOL>valid_length.reshape((-<NUM_LIT:1>, <NUM_LIT:1>)))<EOL>mask = mx.nd.broadcast_axes(mx.nd.expand_dims(mask, axis=<NUM_LIT:1>), axis=<NUM_LIT:1>, size=length)<EOL>if states is None:<EOL><INDENT>states = [mask]<EOL><DEDENT>else:<EOL><INDENT>states.append(mask)<EOL><DEDENT><DEDENT>if self._scale_embed:<EOL><INDENT>inputs = inputs * math.sqrt(inputs.shape[-<NUM_LIT:1>])<EOL><DEDENT>steps = mx.nd.arange(length, ctx=inputs.context)<EOL>if states is None:<EOL><INDENT>states = [steps]<EOL><DEDENT>else:<EOL><INDENT>states.append(steps)<EOL><DEDENT>if valid_length is not None:<EOL><INDENT>step_output, additional_outputs =super(BaseTransformerEncoder, self).forward(inputs, states, valid_length)<EOL><DEDENT>else:<EOL><INDENT>step_output, additional_outputs =super(BaseTransformerEncoder, self).forward(inputs, states)<EOL><DEDENT>return step_output, additional_outputs<EOL>
Parameters
----------
inputs : NDArray, Shape(batch_size, length, C_in)
states : list of NDArray
valid_length : NDArray
steps : NDArray
    Stores value [0, 1, ..., length].
    It is used for lookup in positional encoding matrix

Returns
-------
outputs : NDArray
    The output of the encoder. Shape is (batch_size, length, C_out)
additional_outputs : list
    Either be an empty list or contains the attention weights in this step.
    The attention weights will have shape (batch_size, length, length) or
    (batch_size, num_heads, length, length)