Merge pull request #341 from FluxML/refactor

Refactor and makes internal functions as internal APIs

Author: yuehhua

Project.toml

@@ -1,7 +1,7 @@
 name = "GeometricFlux"
 uuid = "7e08b658-56d3-11e9-2997-919d5b31e4ea"
 authors = ["Yueh-Hua Tu <[email protected]>"]
-version = "0.13.6"
+version = "0.14.0"
 
 [deps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"

src/layers/gn.jl

@@ -75,9 +75,9 @@ See also [`update_edge`](@ref), [`update_vertex`](@ref), [`update_global`](@ref)
 update_batch_edge(gn::GraphNet, el::NamedTuple, E, V, u) =
     update_edge(
         gn,
-        _gather(E, el.es),
-        _gather(V, el.xs),
-        _gather(V, el.nbrs),
+        gather(E, el.es),
+        gather(V, el.xs),
+        gather(V, el.nbrs),
         u
     )
 
@@ -116,17 +116,11 @@ See also [`update_edge`](@ref), [`update_vertex`](@ref), [`update_global`](@ref)
 [`update_batch_edge`](@ref), [`update_batch_vertex`](@ref), [`aggregate_edges`](@ref),
 [`aggregate_vertices`](@ref).
 """
-function aggregate_neighbors(::GraphNet, el::NamedTuple, aggr, E)
-    batch_size = size(E)[end]
-    dstsize = (size(E, 1), el.N, batch_size)
-    xs = batched_index(el.xs, batch_size)
-    return _scatter(aggr, E, xs, dstsize)
-end
-
-aggregate_neighbors(::GraphNet, el::NamedTuple, aggr, E::AbstractMatrix) = _scatter(aggr, E, el.xs)
+aggregate_neighbors(::GraphNet, el::NamedTuple, aggr, E) = scatter(aggr, E, el.xs, el.N)
+aggregate_neighbors(::GraphNet, el::NamedTuple, aggr, E::AbstractMatrix) = scatter(aggr, E, el.xs)
 
-@inline aggregate_neighbors(::GraphNet, ::NamedTuple, ::Nothing, E) = nothing
-@inline aggregate_neighbors(::GraphNet, ::NamedTuple, ::Nothing, ::AbstractMatrix) = nothing
+aggregate_neighbors(::GraphNet, ::NamedTuple, ::Nothing, E) = nothing
+aggregate_neighbors(::GraphNet, ::NamedTuple, ::Nothing, ::AbstractMatrix) = nothing
 
 """
     aggregate_edges(gn, aggr, E)

src/layers/graph_conv.jl

@@ -38,7 +38,7 @@ end
 @functor GCNConv
 
 function (l::GCNConv)(Ã::AbstractMatrix, X::AbstractArray)
-    z = _matmul(l.weight, _matmul(X, Ã))
+    z = matmul(l.weight, matmul(X, Ã))
     return l.σ.(z .+ l.bias)
 end
 
@@ -113,12 +113,12 @@ Flux.trainable(l::ChebConv) = (l.weight, l.bias)
 
 function (l::ChebConv)(L̃::AbstractMatrix, X::AbstractArray)
     Z_prev = X
-    Z = _matmul(X, L̃)
-    Y = _matmul(view(l.weight,:,:,1), Z_prev)
-    Y += _matmul(view(l.weight,:,:,2), Z)
+    Z = matmul(X, L̃)
+    Y = matmul(view(l.weight,:,:,1), Z_prev)
+    Y += matmul(view(l.weight,:,:,2), Z)
     for k = 3:l.k
-        Z, Z_prev = 2 .* _matmul(Z, L̃) .- Z_prev, Z
-        Y += _matmul(view(l.weight,:,:,k), Z)
+        Z, Z_prev = 2 .* matmul(Z, L̃) .- Z_prev, Z
+        Y += matmul(view(l.weight,:,:,k), Z)
     end
     return l.σ.(Y .+ l.bias)
 end
@@ -203,9 +203,9 @@ end
 
 Flux.trainable(l::GraphConv) = (l.weight1, l.weight2, l.bias)
 
-message(gc::GraphConv, x_i, x_j::AbstractArray, e_ij) = _matmul(gc.weight2, x_j)
+message(gc::GraphConv, x_i, x_j::AbstractArray, e_ij) = matmul(gc.weight2, x_j)
 
-update(gc::GraphConv, m::AbstractArray, x::AbstractArray) = gc.σ.(_matmul(gc.weight1, x) .+ m .+ gc.bias)
+update(gc::GraphConv, m::AbstractArray, x::AbstractArray) = gc.σ.(matmul(gc.weight1, x) .+ m .+ gc.bias)
 
 # For variable graph
 function (l::GraphConv)(fg::AbstractFeaturedGraph)
@@ -299,7 +299,7 @@ function update_batch_edge(gat::GATConv, el::NamedTuple, E, X::AbstractMatrix, u
 end
 
 function update_batch_edge(gat::GATConv, el::NamedTuple, E, X::AbstractArray, u)
-    Xi, Xj = _gather(X, el.xs), _gather(X, el.nbrs)
+    Xi, Xj = gather(X, el.xs), gather(X, el.nbrs)
     _, nb, bch_sz = size(Xj)
     heads = gat.heads
     Q = reshape(NNlib.batched_mul(gat.weight, Xi), :, heads, nb, bch_sz)  # dims: (out, heads, nb, bch_sz)
@@ -429,7 +429,7 @@ function update_batch_edge(gat::GATv2Conv, el::NamedTuple, E, X::AbstractMatrix,
 end
 
 function update_batch_edge(gat::GATv2Conv, el::NamedTuple, E, X::AbstractArray, u)
-    Xi, Xj = _gather(X, el.xs), _gather(X, el.nbrs)
+    Xi, Xj = gather(X, el.xs), gather(X, el.nbrs)
     _, nb, bch_sz = size(Xj)
     heads = gat.heads
     Q = reshape(NNlib.batched_mul(gat.wi, Xi) .+ gat.biasi, :, heads, nb, bch_sz)  # dims: (out, heads, nb, bch_sz)
@@ -556,7 +556,7 @@ function (l::GatedGraphConv)(el::NamedTuple, H::AbstractArray{T}) where {T<:Real
         H = vcat(H, Hpad)
     end
     for i = 1:l.num_layers
-        M = _matmul(view(l.weight, :, :, i), H)
+        M = matmul(view(l.weight, :, :, i), H)
         _, M = propagate(l, el, nothing, M, nothing, l.aggr, nothing, nothing)
         H, _ = l.gru(H, M)
     end
@@ -822,20 +822,17 @@ end
 message(l::SAGEConv, x_i, x_j::AbstractArray, e) = l.proj(x_j)
 
 function aggregate_neighbors(l::SAGEConv, el::NamedTuple, aggr, E)
-    batch_size = size(E)[end]
     sample_idx = sample_node_index(E, l.num_sample; dims=2)
-    idx = ntuple(i -> (i == 2) ? sample_idx : Colon(), ndims(E))
-    dstsize = (size(E, 1), el.N, batch_size)  # ensure outcome has the same dimension as x in update
-    xs = batched_index(el.xs[sample_idx], batch_size)
-    Ē = _scatter(aggr, E[idx...], xs, dstsize)
+    indexed_E = selectdim(E, 2, sample_idx)
+    Ē = scatter(aggr, indexed_E, el.xs[sample_idx], el.N)
     return Ē
 end
 
 function aggregate_neighbors(l::SAGEConv, el::NamedTuple, aggr, E::AbstractMatrix)
     sample_idx = sample_node_index(E, l.num_sample; dims=2)
-    idx = ntuple(i -> (i == 2) ? sample_idx : Colon(), ndims(E))
+    indexed_E = selectdim(E, 2, sample_idx)
     dstsize = (size(E, 1), el.N)  # ensure outcome has the same dimension as x in update
-    Ē = _scatter(aggr, E[idx...], el.xs[sample_idx], dstsize)
+    Ē = NNlib.scatter(aggr, indexed_E, el.xs[sample_idx]; dstsize=dstsize)
     return Ē
 end
 
@@ -844,13 +841,13 @@ aggregate_neighbors(::SAGEConv, el::NamedTuple, lstm::Flux.LSTMCell, E::Abstract
 
 function aggregate_neighbors(::SAGEConv, el::NamedTuple, lstm::Flux.LSTMCell, E::AbstractMatrix)
     sample_idx = sample_node_index(E, el.N; dims=2)
-    idx = ntuple(i -> (i == 2) ? sample_idx : Colon(), ndims(E))
-    state, Ē = lstm(lstm.state0, E[idx...])
+    indexed_E = selectdim(E, 2, sample_idx)
+    state, Ē = lstm(lstm.state0, indexed_E)
     return Ē
 end
 
 function update(l::SAGEConv, m::AbstractArray, x::AbstractArray)
-    y = l.σ.(_matmul(l.weight1, x) + _matmul(l.weight2, m) .+ l.bias)
+    y = l.σ.(matmul(l.weight1, x) + matmul(l.weight2, m) .+ l.bias)
     l.normalize && (y = l2normalize(y; dims=2))  # across all nodes
     return y
 end

src/layers/group_conv.jl

@@ -104,26 +104,20 @@ function Base.show(io::IO, l::EEquivGraphConv)
     print(io, ")")
 end
 
-function aggregate_neighbors(::EEquivGraphConv, el::NamedTuple, aggr, E)
-    batch_size = size(E)[end]
-    dstsize = (size(E, 1), el.N, batch_size)
-    xs = batched_index(el.xs, batch_size)
-    return _scatter(aggr, E, xs, dstsize)
-end
-
-aggregate_neighbors(::EEquivGraphConv, el::NamedTuple, aggr, E::AbstractMatrix) = _scatter(aggr, E, el.xs)
+aggregate_neighbors(::EEquivGraphConv, el::NamedTuple, aggr, E) = scatter(aggr, E, el.xs, el.N)
+aggregate_neighbors(::EEquivGraphConv, el::NamedTuple, aggr, E::AbstractMatrix) = scatter(aggr, E, el.xs)
 
-@inline aggregate_neighbors(::EEquivGraphConv, ::NamedTuple, ::Nothing, E) = nothing
-@inline aggregate_neighbors(::EEquivGraphConv, ::NamedTuple, ::Nothing, ::AbstractMatrix) = nothing
+aggregate_neighbors(::EEquivGraphConv, ::NamedTuple, ::Nothing, E) = nothing
+aggregate_neighbors(::EEquivGraphConv, ::NamedTuple, ::Nothing, ::AbstractMatrix) = nothing
 
 propagate(l::EEquivGraphConv, sg::SparseGraph, E, V, X, aggr) =
     propagate(l, GraphSignals.to_namedtuple(sg), E, V, X, aggr)
 
 function propagate(l::EEquivGraphConv, el::NamedTuple, E, V, X, aggr)
     E = message(
-        l, _gather(V, el.xs), _gather(V, el.nbrs),
-        _gather(X, el.xs), _gather(X, el.nbrs),
-        _gather(E, el.es)
+        l, gather(V, el.xs), gather(V, el.nbrs),
+        gather(X, el.xs), gather(X, el.nbrs),
+        gather(E, el.es)
         )
     X = positional_encode(l.pe, el, X, E)
     Ē = aggregate_neighbors(l, el, aggr, E)

src/layers/positional.jl

@@ -275,20 +275,13 @@ update_edge(l::GatedGCNLSPEConv, h_i, h_j, e_ij) = σ.(l.B1(h_i) + l.B2(h_j) + l
 
 function normalize_η(l::GatedGCNLSPEConv, el::NamedTuple, η̂)
     summed_η = aggregate_neighbors(l, el, +, η̂)
-    return η̂ ./ _gather(summed_η .+ 1f-6, el.xs)
+    return η̂ ./ gather(summed_η .+ 1f-6, el.xs)
 end
 
 message_vertex(l::GatedGCNLSPEConv, h_j, p_j, η_ij) = η_ij .* l.A2(vcat(h_j, p_j))
 
-function aggregate_neighbors(l::GatedGCNLSPEConv, el::NamedTuple, aggr, E)
-    batch_size = size(E)[end]
-    dstsize = (size(E, 1), el.N, batch_size)
-    xs = batched_index(el.xs, batch_size)
-    return _scatter(aggr, E, xs, dstsize)
-end
-
-aggregate_neighbors(l::GatedGCNLSPEConv, el::NamedTuple, aggr, E::AbstractMatrix) =
-    _scatter(aggr, E, el.xs)
+aggregate_neighbors(l::GatedGCNLSPEConv, el::NamedTuple, aggr, E) = scatter(aggr, E, el.xs, el.N)
+aggregate_neighbors(l::GatedGCNLSPEConv, el::NamedTuple, aggr, E::AbstractMatrix) = scatter(aggr, E, el.xs)
 
 update_vertex(l::GatedGCNLSPEConv, m, h, p) = l.σ.(l.A1(vcat(h, p)) + m)
 
@@ -300,11 +293,11 @@ propagate(l::GatedGCNLSPEConv, sg::SparseGraph, E, H, X) =
     propagate(l, GraphSignals.to_namedtuple(sg), E, H, X)
 
 function propagate(l::GatedGCNLSPEConv, el::NamedTuple, E, H, X)
-    e_ij = _gather(E, el.es)
-    h_i = _gather(H, el.xs)
-    h_j = _gather(H, el.nbrs)
-    p_i = _gather(X, el.xs)
-    p_j = _gather(X, el.nbrs)
+    e_ij = gather(E, el.es)
+    h_i = gather(H, el.xs)
+    h_j = gather(H, el.nbrs)
+    p_i = gather(X, el.xs)
+    p_j = gather(X, el.nbrs)
 
     η̂ = update_edge(l, h_i, h_j, e_ij)
     Ê = l.σ.(η̂)

src/operation.jl

@@ -1,13 +1,19 @@
-_gather(::Nothing, idx) = nothing
-_gather(A::Fill{T,2,Axes}, idx) where {T,Axes} = fill(A.value, A.axes[1], length(idx))
-_gather(A::AbstractMatrix, idx) = NNlib.gather(A, idx)
-_gather(A::AbstractArray, idx) = NNlib.gather(A, batched_index(idx, size(A)[end]))
+gather(::Nothing, idx) = nothing
+gather(A::Fill{T,2,Axes}, idx) where {T,Axes} = fill(A.value, A.axes[1], length(idx))
+gather(A::AbstractMatrix, idx) = NNlib.gather(A, idx)
+gather(A::AbstractArray, idx) = NNlib.gather(A, batched_index(idx, size(A)[end]))
 
-_scatter(aggr, E, xs::AbstractArray) = NNlib.scatter(aggr, E, xs)
-_scatter(aggr, E, xs::AbstractArray, dstsize) = NNlib.scatter(aggr, E, xs; dstsize=dstsize)
+scatter(aggr, E, xs::AbstractArray) = NNlib.scatter(aggr, E, xs)
 
-_matmul(A::AbstractMatrix, B::AbstractMatrix) = A * B
-_matmul(A::AbstractArray, B::AbstractArray) = NNlib.batched_mul(A, B)
+function scatter(aggr, E, xs::AbstractArray, N::Int)
+    dim, batch_size = size(E)[1], size(E)[end]
+    dstsize = (dim, N, batch_size)
+    batched_xs = batched_index(xs, batch_size)
+    return NNlib.scatter(aggr, E, batched_xs; dstsize=dstsize)
+end
+
+matmul(A::AbstractMatrix, B::AbstractMatrix) = A * B
+matmul(A::AbstractArray, B::AbstractArray) = NNlib.batched_mul(A, B)
 
 function batched_index(idx::AbstractVector, batch_size::Integer)
     b = copyto!(similar(idx, 1, batch_size), collect(1:batch_size))
@@ -36,29 +42,20 @@ end
 function indexed_softmax(x::AbstractArray, xs, N; dims=1)
     y = copy(x)
     for i in 1:N
-        idx = ntuple(j -> (j == dims) ? (xs .== i) : Colon(), ndims(y))
-        NNlib.softmax!(view(y, idx...); dims)
+        indexed_y = selectdim(y, dims, (xs .== i))
+        NNlib.softmax!(indexed_y; dims)
     end
     return y
 end
 
 function ∇indexed_softmax(dy::AbstractArray{T}, y::AbstractArray{S}, xs, N; dims=1) where {T,S}
-    dx = if NNlib.within_grad()
-        tmp = dy .* y
-        for i in 1:N
-            idx = ntuple(j -> (j == dims) ? (xs .== i) : Colon(), ndims(y))
-            tmp[idx...] .= tmp[idx...] .- y[idx...] .* sum(tmp[idx...]; dims)
-        end
-        tmp
-    else
-        out = similar(y, promote_type(T,S))
-        out .= dy .* y
-        for i in 1:N
-            idx = ntuple(j -> (j == dims) ? (xs .== i) : Colon(), ndims(y))
-            out[idx...] .= out[idx...] .- y[idx...] .* sum(out[idx...]; dims)
-        end
-        out
+    out = NNlib.∇softmax_data(dy, y; dims)
+    for i in 1:N
+        indexed_y = selectdim(y, dims, (xs .== i))
+        indexed_out = selectdim(out, dims, (xs .== i))
+        indexed_out .= indexed_out .- indexed_y .* sum(indexed_out; dims)
     end
+    return out
 end
 
 function ChainRulesCore.rrule(::typeof(indexed_softmax), x, xs, N; dims=1)