Merge pull request #57 from yuehhua/fix

Ignore gradient of generate_cluster

Author: yuehhua

src/GeometricFlux.jl

@@ -3,14 +3,15 @@ module GeometricFlux
 using Statistics: mean
 using StaticArrays: StaticArray
 using SparseArrays: SparseMatrixCSC
-using LinearAlgebra: I, issymmetric, diagm, eigmax, norm, Adjoint
+using LinearAlgebra: I, issymmetric, diagm, eigmax, norm, Adjoint, Diagonal
 
 using Requires
 using DataStructures: DefaultDict
 using Flux
 using Flux: glorot_uniform, leakyrelu, GRUCell
 using Flux: @functor
 using LightGraphs
+using Zygote
 using ZygoteRules
 using FillArrays: Fill
 

src/layers/meta.jl

@@ -34,7 +34,8 @@ function propagate(meta::T; adjl=adjlist(meta), kwargs...) where {T<:Meta}
     newE = update_edge(meta; gi=gi, kwargs...)
 
     if haskey(kwargs, :naggr)
-        Ē = aggregate_neighbors(meta, kwargs[:naggr]; E=newE, cluster=generate_cluster(newE, gi))
+        cluster = generate_cluster(newE, gi)
+        Ē = aggregate_neighbors(meta, kwargs[:naggr]; E=newE, cluster=cluster)
         kwargs = (kwargs..., Ē=Ē)
     end
 
@@ -54,7 +55,7 @@ function propagate(meta::T; adjl=adjlist(meta), kwargs...) where {T<:Meta}
     (newE, newV, new_u)
 end
 
-function generate_cluster(M::AbstractArray{T,N}, gi::GraphInfo) where {T,N}
+Zygote.@nograd function generate_cluster(M::AbstractArray{T,N}, gi::GraphInfo) where {T,N}
     cluster = similar(M, Int, gi.E)
     @inbounds for i = 1:gi.V
         j = gi.edge_idx[i]

src/layers/msgpass.jl

@@ -41,7 +41,8 @@ function propagate(mp::T; aggr::Symbol=:add, adjl=adjlist(mp), kwargs...) where
     M = update_edge(mp; gi=gi, kwargs...)
 
     # aggregate function
-    M = aggregate_neighbors(mp, aggr; M=M, cluster=generate_cluster(M, gi))
+    cluster = generate_cluster(M, gi)
+    M = aggregate_neighbors(mp, aggr; M=M, cluster=cluster)
 
     # update function
     Y = update_vertex(mp; M=M, kwargs...)

src/operations/linalg.jl

@@ -1,5 +1,6 @@
 ## Linear algebra API for adjacency matrix
-using LinearAlgebra
+
+Zygote.@nograd issymmetric
 
 function adjacency_matrix(adj::AbstractMatrix, T::DataType=eltype(adj))
     m, n = size(adj)
@@ -124,9 +125,9 @@ function normalized_laplacian(adj::AbstractMatrix, T::DataType=eltype(adj); self
 end
 
 @doc raw"""
-    scaled_laplacian(adj::AbstractMatrix[, T::DataType]) 
+    scaled_laplacian(adj::AbstractMatrix[, T::DataType])
 
-Scaled Laplacien matrix of graph `g`, 
+Scaled Laplacien matrix of graph `g`,
 defined as ``\hat{L} = \frac{2}{\lambda_{max}} L - I`` where ``L`` is the normalized Laplacian matrix.
 
 # Arguments