New add_self_loops(g) for hetero graphs (#402)

* new add self loops for hetero graphs

* added support for sgconv

* concise

* newline

* Update transform.jl

* add self loop

* change docs

* docs

* docs

Author: rbSparky

src/GNNGraphs/transform.jl

@@ -40,6 +40,7 @@ end
 
 """
     add_self_loops(g::GNNHeteroGraph, edge_t::EType)
+    add_self_loops(g::GNNHeteroGraph)
 
 If the source node type is the same as the destination node type in `edge_t`,
 return a graph with the same features as `g` but also add self-loops 
@@ -51,7 +52,10 @@ a second set of self-loops of the same type.
 If the graph has edge weights for edges of type `edge_t`, the new edges will have weight 1.
 
 If no edges of type `edge_t` exist, or all existing edges have no weight, 
-then all new self-loops will have no weight.
+then all new self loops will have no weight.
+
+If `edge_t` is not passed as argument, for the entire graph self-loop is added to each node for every edge type in the graph where the source and destination node types are the same. 
+This iterates over all edge types present in the graph, applying the self-loop addition logic to each applicable edge type.
 """
 function add_self_loops(g::GNNHeteroGraph{Tuple{T, T, V}}, edge_t::EType) where {T <: AbstractVector{<:Integer}, V}
     function get_edge_weight_nullable(g::GNNHeteroGraph{<:COO_T}, edge_t::EType)
@@ -99,6 +103,13 @@ function add_self_loops(g::GNNHeteroGraph{Tuple{T, T, V}}, edge_t::EType) where
              ntypes, etypes)
 end
 
+function add_self_loops(g::GNNHeteroGraph)
+    for edge_t in keys(g.graph)
+        g = add_self_loops(g, edge_t)
+    end
+    return g
+end
+
 """
     remove_self_loops(g::GNNGraph)
 

src/layers/conv.jl

@@ -419,11 +419,7 @@ function (l::GATConv)(g::AbstractGNNGraph, x,
 
     if l.add_self_loops
         @assert e===nothing "Using edge features and setting add_self_loops=true at the same time is not yet supported."
-       if g isa GNNHeteroGraph
-           g = add_self_loops(g,  g.etypes[1])
-       else
-          g = add_self_loops(g)
-       end
+        g = add_self_loops(g)
     end
 
     _, chout = l.channel
@@ -1530,7 +1526,7 @@ function (l::SGConv)(g::AbstractGNNGraph, x,
     end
 
     if l.add_self_loops
-        g = g isa GNNHeteroGraph ? add_self_loops(g, edge_t) : add_self_loops(g)
+        g = add_self_loops(g)
         if edge_weight !== nothing
             edge_weight = [edge_weight; fill!(similar(edge_weight, g.num_nodes), 1)]
             @assert length(edge_weight) == g.num_edges