Implement $D^{-1/2}AD^{-1/2}H$

[kirovreporting]

I am trying to implement D^{-1/2}AD^{-1/2}H using update_all. So far the message and reduce functions I were only able to come up with are

mess_f = fn.v_mul_u('deg', 'h', 'm')
reduce_f = fn.sum('m', 'h')
g_water.update_all(mess_f, reduce_f)

In which

g_water.ndata['h'] = torch.tensor([1,0,0,0,0]).float()
g_water.ndata['deg'] = 1/torch.sqrt(g_water.in_degrees())

I think what I am implementing is D^{-1/2}AH. A naive thing I thought of is mess_f = fn.u_mul_v(fn.v_mul_u('deg', 'h', 'm'), 'deg', 'm'), but of course it wouldn’t work with DGL.

How do I implement this?

[minjie]

Hi @kirovreporting , first of all, please try out the recently released sparse API, which is designed exactly for such case. The GCN tutorial demonstrates exactly how to implement D^{-1/2}AD^{-1/2}H.

The trick to implement it in the message passing interface is to realize that it can be seen as a symmetric normalization – the inverse square root of the degree vector need to be multiplied before and after the message passing operation. Below is an illustration:

g = ...  # your graph
deg_invsqrt = 1 / torch.sqrt(g.in_degrees())  # <-- BTW, watch for zero degrees
g.ndata['h'] = torch.tensor(...) * deg_invsqrt
g.update_all(fn.copy_u('h', 'm'), fn.sum('m', 'h'))
h_new = g.ndata['h'] * deg_invsqrt