def forward(self, graph, feat, distances):
        """

        Description
        -----------
        Apply the atomic convolution layer.

        Parameters
        ----------
        graph : DGLGraph
            Topology based on which message passing is performed.
        feat : Float32 tensor of shape :math:`(V, 1)`
            Initial node features, which are atomic numbers in the paper.
            :math:`V` for the number of nodes.
        distances : Float32 tensor of shape :math:`(E, 1)`
            Distance between end nodes of edges. E for the number of edges.

        Returns
        -------
        Float32 tensor of shape :math:`(V, K * T)`
            Updated node representations. :math:`V` for the number of nodes, :math:`K` for the
            number of radial filters, and :math:`T` for the number of types of atomic numbers.
        """
        with graph.local_scope():
            radial_pooled_values = self.radial_pooling(distances).to(
                feat
            )  # (K, E, 1)
            if self.features_to_use is not None:
                feat = (feat == self.features_to_use).to(feat)  # (V, T)
            graph.ndata["hv"] = feat
            graph.edata["he"] = radial_pooled_values.transpose(1, 0).squeeze(
                -1
            )  # (E, K)
            graph.update_all(msg_func, reduce_func)

            return graph.ndata["hv_new"].view(
                graph.num_nodes(), -1
            )  # (V, K * T)


        Description
        -----------
        Apply the atomic convolution layer.

        Parameters
        ----------
        graph : DGLGraph
            Topology based on which message passing is performed.
        feat : Float32 tensor of shape :math:`(V, 1)`
            Initial node features, which are atomic numbers in the paper.
            :math:`V` for the number of nodes.
        distances : Float32 tensor of shape :math:`(E, 1)`
            Distance between end nodes of edges. E for the number of edges.

        Returns
        -------
        Float32 tensor of shape :math:`(V, K * T)`
            Updated node representations. :math:`V` for the number of nodes, :math:`K` for the
            number of radial filters, and :math:`T` for the number of types of atomic numbers.
        
python/dgl/nn/pytorch/conv/atomicconv.py