def METHOD_NAME(self, edges):
    """message function for EGNN"""
    # concat features for edge mlp
    if self.edge_feat_size > 0:
        f = torch.cat(
            [
                edges.src["h"],
                edges.dst["h"],
                edges.data["radial"],
                edges.data["a"],
            ],
            dim=-1,
        )
    else:
        f = torch.cat(
            [edges.src["h"], edges.dst["h"], edges.data["radial"]], dim=-1
        )
    msg_h = self.edge_mlp(f)
    msg_x = self.coord_mlp(msg_h) * edges.data["x_diff"]
    return {"msg_x": msg_x, "msg_h": msg_h}