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
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.
        
python/dgl/nn/mxnet/glob.py