code: def get_context_and_attention_probs(values: mx.sym.Symbol,
                                      length: mx.sym.Symbol,
                                      logits: mx.sym.Symbol,
                                      dtype: str) -> Tuple[mx.sym.Symbol, mx.sym.Symbol]:
      
      # masks attention scores according to sequence length.
      # (batch_size, seq_len, 1)
      logits = mx.sym.SequenceMask(data=logits,
                                   axis=1,
                                   use_sequence_length=True,
                                   sequence_length=length,
                                   value=-C.LARGE_VALUES[dtype])
  
      # (batch_size, seq_len, 1)
      probs = mx.sym.softmax(logits, axis=1, name='attention_softmax')
  
      # batch_dot: (batch, M, K) X (batch, K, N) –> (batch, M, N).
      # (batch_size, seq_len, num_hidden) X (batch_size, seq_len, 1) -> (batch_size, num_hidden, 1)
      context = mx.sym.batch_dot(lhs=values, rhs=probs, transpose_a=True)
      # (batch_size, encoder_num_hidden, 1)-> (batch_size, encoder_num_hidden)
      context = mx.sym.reshape(data=context, shape=(0, 0))
      probs = mx.sym.reshape(data=probs, shape=(0, 0))
  
      return context, probs
docs: Returns context vector and attention probabilities
      via a weighted sum over values.
  
      :param values: Shape: (batch_size, seq_len, encoder_num_hidden).
      :param length: Shape: (batch_size,).
      :param logits: Shape: (batch_size, seq_len, 1).
      :param dtype: data type.
      :return: context: (batch_size, encoder_num_hidden), attention_probs: (batch_size, seq_len).