class_docstr: Set2Set operator from `Order Matters: Sequence to sequence for sets
  <https://arxiv.org/pdf/1511.06391.pdf>`__
  
  For each individual graph in the batch, set2set computes
  
  .. math::
      q_t &= \mathrm{LSTM} (q^*_{t-1})
  
      \alpha_{i,t} &= \mathrm{softmax}(x_i \cdot q_t)
  
      r_t &= \sum_{i=1}^N \alpha_{i,t} x_i
  
      q^*_t &= q_t \Vert r_t
  
  for this graph.
  
  Parameters
  ----------
  input_dim : int
      Size of each input sample
  n_iters : int
      Number of iterations.
  n_layers : int
      Number of recurrent layers.
docstr: Compute set2set pooling.
  
  Parameters
  ----------
  graph : DGLGraph
      The graph.
  feat : mxnet.NDArray
      The input node feature with shape :math:`(N, D)` where
      :math:`N` is the number of nodes in the graph.
  
  Returns
  -------
  mxnet.NDArray
      The output feature with shape :math:`(B, D)`, where
      :math:`B` refers to the batch size.
docstr_tok: ["Compute","set2set","pooling",".","Parameters","--","--","--","--","--","graph",":","DGLGraph","The","graph",".","feat",":","mxnet.NDArray","The","input","node","feature","with","shape",":","math",":","`","(","N",",","D",")","`","where",":","math",":","`","N","`","is","the","number","of","nodes","in","the","graph",".","Returns","--","--","--","-","mxnet.NDArray","The","output","feature","with","shape",":","math",":","`","(","B",",","D",")","`",",","where",":","math",":","`","B","`","refers","to","the","batch","size","."]
code: def forward(self, graph, feat):
      r"""Compute set2set pooling.
  
      Parameters
      ----------
      graph : DGLGraph
          The graph.
      feat : mxnet.NDArray
          The input node feature with shape :math:`(N, D)` where
          :math:`N` is the number of nodes in the graph.
  
      Returns
      -------
      mxnet.NDArray
          The output feature with shape :math:`(B, D)`, where
          :math:`B` refers to the batch size.
      """
      with graph.local_scope():
          batch_size = graph.batch_size
  
          h = (
              nd.zeros(
                  (self.n_layers, batch_size, self.input_dim),
                  ctx=feat.context,
              ),
              nd.zeros(
                  (self.n_layers, batch_size, self.input_dim),
                  ctx=feat.context,
              ),
          )
          q_star = nd.zeros(
              (batch_size, self.output_dim), ctx=feat.context
          )
  
          for _ in range(self.n_iters):
              q, h = self.lstm(q_star.expand_dims(axis=0), h)
              q = q.reshape((batch_size, self.input_dim))
              e = (feat * broadcast_nodes(graph, q)).sum(
                  axis=-1, keepdims=True
              )
              graph.ndata["e"] = e
              alpha = softmax_nodes(graph, "e")
              graph.ndata["r"] = feat * alpha
              readout = sum_nodes(graph, "r")
              q_star = nd.concat(q, readout, dim=-1)
  
          return q_star
  
imports: from mxnet import gluon, nd
  from mxnet.gluon import nn
  from ...readout import broadcast_nodes, max_nodes, mean_nodes, softmax_nodes, sum_nodes, topk_nodes
label_desc: Use image node_id 2 for calling the Set2Set obj's underlying member method code with example usage: obj.forward(graph, feat) and returns: q_star