Utility Functions
Index
GraphNeuralNetworks.broadcast_edges
GraphNeuralNetworks.broadcast_nodes
GraphNeuralNetworks.reduce_edges
GraphNeuralNetworks.reduce_nodes
GraphNeuralNetworks.softmax_edge_neighbors
GraphNeuralNetworks.softmax_edges
GraphNeuralNetworks.softmax_nodes
NNlib.gather
NNlib.gather!
NNlib.scatter
NNlib.scatter!
Docs
Graph-wise operations
GraphNeuralNetworks.reduce_nodes
— Functionreduce_nodes(aggr, g, x)
For a batched graph g
, return the graph-wise aggregation of the node features x
. The aggregation operator aggr
can be +
, mean
, max
, or min
. The returned array will have last dimension g.num_graphs
.
GraphNeuralNetworks.reduce_edges
— Functionreduce_edges(aggr, g, e)
For a batched graph g
, return the graph-wise aggregation of the edge features e
. The aggregation operator aggr
can be +
, mean
, max
, or min
. The returned array will have last dimension g.num_graphs
.
GraphNeuralNetworks.softmax_nodes
— Functionsoftmax_nodes(g, x)
Graph-wise softmax of the node features x
.
GraphNeuralNetworks.softmax_edges
— Functionsoftmax_edges(g, e)
Graph-wise softmax of the edge features e
.
GraphNeuralNetworks.broadcast_nodes
— Functionbroadcast_nodes(g, x)
Graph-wise broadcast array x
of size (*, g.num_graphs)
to size (*, g.num_nodes)
.
GraphNeuralNetworks.broadcast_edges
— Functionbroadcast_edges(g, x)
Graph-wise broadcast array x
of size (*, g.num_graphs)
to size (*, g.num_edges)
.
Neighborhood operations
GraphNeuralNetworks.softmax_edge_neighbors
— Functionsoftmax_edge_neighbors(g, e)
Softmax over each node's neighborhood of the edge features e
.
\[\mathbf{e}'_{j\to i} = \frac{e^{\mathbf{e}_{j\to i}}} {\sum_{j'\in N(i)} e^{\mathbf{e}_{j'\to i}}}.\]
NNlib
Primitive functions implemented in NNlib.jl.
NNlib.gather!
— FunctionNNlib.gather!(dst, src, idx)
Reverse operation of scatter!
. Gathers data from source src
and writes it in destination dst
according to the index array idx
. For each k
in CartesianIndices(idx)
, assign values to dst
according to
dst[:, ... , k] .= src[:, ... , idx[k]...]
Notice that if idx
is a vector containing integers, and both dst
and src
are matrices, previous expression simplifies to
dst[:, k] .= src[:, idx[k]]
and k
will run over 1:length(idx)
.
The elements of idx
can be integers or integer tuples and may be repeated. A single src
column can end up being copied into zero, one, or multiple dst
columns.
See gather
for an allocating version.
NNlib.gather
— FunctionNNlib.gather(src, idx) -> dst
Reverse operation of scatter
. Gathers data from source src
and writes it in a destination dst
according to the index array idx
. For each k
in CartesianIndices(idx)
, assign values to dst
according to
dst[:, ... , k] .= src[:, ... , idx[k]...]
Notice that if idx
is a vector containing integers and src
is a matrix, previous expression simplifies to
dst[:, k] .= src[:, idx[k]]
and k
will run over 1:length(idx)
.
The elements of idx
can be integers or integer tuples and may be repeated. A single src
column can end up being copied into zero, one, or multiple dst
columns.
See gather!
for an in-place version.
Examples
julia> NNlib.gather([1,20,300,4000], [2,4,2])
3-element Vector{Int64}:
20
4000
20
julia> NNlib.gather([1 2 3; 4 5 6], [1,3,1,3,1])
2×5 Matrix{Int64}:
1 3 1 3 1
4 6 4 6 4
NNlib.scatter!
— FunctionNNlib.scatter!(op, dst, src, idx)
Scatter operation, which writes data in src
into dst
at locations idx
. A binary reduction operator op
is applied during the scatter. For each index k
in idx
, accumulates values in dst
according to
dst[:, ..., idx[k]...] = (op).(dst[:, ..., idx[k]...], src[:, ..., k...])
Arguments
op
: Operations to be applied ondst
andsrc
, e.g.+
,-
,*
,/
,max
,min
andmean
.dst
: The destination forsrc
to aggregate to. This argument will be mutated.src
: The source data for aggregating.idx
: The mapping for aggregation from source (index) to destination (value). Theidx
array can contain either integers or tuples.
Examples
julia> NNlib.scatter!(+, ones(3), [10,100], [1,3])
3-element Vector{Float64}:
11.0
1.0
101.0
julia> NNlib.scatter!(*, fill(0.5, 2, 4), [1 10; 100 1000], [3,2])
2×4 Matrix{Float64}:
0.5 5.0 0.5 0.5
0.5 500.0 50.0 0.5
NNlib.scatter
— FunctionNNlib.scatter(op, src, idx; [init, dstsize])
Scatter operation allocating a destination array dst
and calling scatter!(op, dst, src, idx)
on it.
If keyword
init
is provided, it is used to initialize the content ofdst
. Otherwise, the init values is inferred from the reduction operatorop
for some common operators (e.g.init = 0
forop = +
).If
dstsize
is provided, it will be used to define the size of destination array, otherwise it will be inferred bysrc
andidx
.
See scatter!
for full details on how idx
works.
Examples
julia> NNlib.scatter(+, [10,100,1000], [3,1,2])
3-element Vector{Int64}:
100
1000
10
julia> NNlib.scatter(+, [1 2 3 4; 5 6 7 8], [2,1,1,5])
2×5 Matrix{Int64}:
5 1 0 0 4
13 5 0 0 8
julia> NNlib.scatter(*, [10,200,3000], [1,4,2]; init = 10, dstsize = 6)
6-element Vector{Int64}:
100
30000
10
2000
10
10