code: def _attend(self,
                  queries: mx.sym.Symbol,
                  keys: mx.sym.Symbol,
                  values: mx.sym.Symbol,
                  lengths: Optional[mx.sym.Symbol] = None,
                  bias: Optional[mx.sym.Symbol] = None) -> mx.sym.Symbol:
          
          # scale by sqrt(depth_per_head)
          queries = queries * (self.depth_per_head ** -0.5)
  
          # (batch*heads, length, depth/heads)
          queries = split_heads(queries, self.depth_per_head, self.heads)
          keys = split_heads(keys, self.depth_per_head, self.heads)
          values = split_heads(values, self.depth_per_head, self.heads)
          lengths = broadcast_to_heads(lengths, self.heads, ndim=1, fold_heads=True) if lengths is not None else lengths
  
          # (batch*heads, query_max_length, depth_per_head)
          contexts = dot_attention(queries, keys, values,
                                   lengths=lengths, dropout=self.dropout, bias=bias, prefix=self.prefix)
  
          # (batch, query_max_length, depth)
          contexts = combine_heads(contexts, self.depth_per_head, self.heads)
  
          # contexts: (batch, query_max_length, output_depth)
          contexts = mx.sym.FullyConnected(data=contexts,
                                           weight=self.w_h2o,
                                           no_bias=True,
                                           num_hidden=self.depth_out,
                                           flatten=False)
  
          return contexts
docs: Returns context vectors of multi-head dot attention.
  
          :param queries: Query tensor. Shape: (batch_size, query_max_length, depth).
          :param keys: Keys. Shape: (batch_size, memory_max_length, depth).
          :param values: Values. Shape: (batch_size, memory_max_length, depth).
          :param lengths: Optional lengths of keys. Shape: (batch_size,).
          :param bias: Optional 3d bias.
          :return: Context vectors. Shape: (batch_size, query_max_length, output_depth).