Utility Functions

reduce_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.

See also: reduce_edges.

reduce_nodes(aggr, indicator::AbstractVector, x)

Return the graph-wise aggregation of the node features x given the graph indicator indicator. The aggregation operator aggr can be +, mean, max, or min.

See also graph_indicator.

reduce_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.

softmax_nodes(g, x)

Graph-wise softmax of the node features x.

softmax_edges(g, e)

Graph-wise softmax of the edge features e.

broadcast_nodes(g, x)

Graph-wise broadcast array x of size (*, g.num_graphs) to size (*, g.num_nodes).

broadcast_edges(g, x)

Graph-wise broadcast array x of size (*, g.num_graphs) to size (*, g.num_edges).

softmax_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.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(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

gather(src, IJK...)

Convert the tuple of integer vectors IJK to a tuple of CartesianIndex and call gather on it: gather(src, CartesianIndex.(IJK...)).

Examples

julia> src = reshape([1:15;], 3, 5)
3×5 Matrix{Int64}:
 1  4  7  10  13
 2  5  8  11  14
 3  6  9  12  15

julia> NNlib.gather(src, [1, 2], [2, 4])
2-element Vector{Int64}:
  4
 11

NNlib.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...])

See also scatter, gather.

Arguments

• op: Operations to be applied on dst and src, e.g. +, -, *, /, max, min and mean.
• dst: The destination for src 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). The idx 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(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 of dst. Otherwise, the init values is inferred from the reduction operator op for some common operators (e.g. init = 0 for op = +).

• If dstsize is provided, it will be used to define the size of destination array, otherwise it will be inferred by src and idx.

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