GNNLux MEGNetConv drops the state of stateful sublayers #732

GNNLux.MEGNetConv wraps ϕe and ϕv in StatefulLuxLayer{true}, but its forward pass returns the input state unchanged:

function (l::MEGNetConv)(g, x, e, ps, st)
    ϕe = StatefulLuxLayer{true}(l.ϕe, ps.ϕe, _getstate(st, :ϕe))
    ϕv = StatefulLuxLayer{true}(l.ϕv, ps.ϕv, _getstate(st, :ϕv))
    m = (; ϕe, ϕv, aggr=l.aggr)
    return GNNlib.megnet_conv(m, g, x, e), st
end

Any state update performed by a sublayer is therefore dropped. The default constructor only builds Dense chains, which are stateless, but the MEGNetConv(in_dims, out_dims, ϕe, ϕv; aggr) constructor accepts arbitrary layers. With a stateful one, e.g. BatchNorm or Dropout inside ϕe/ϕv, the running statistics (or RNG) never advance in training mode.

The same pattern is already used by EdgeConv, GINConv and, since #728, NNConv:

stnew = (; ϕe = _getstate(ϕe), ϕv = _getstate(ϕv))
return GNNlib.megnet_conv(m, g, x, e), stnew

This was found by reading the code while reviewing #728 and has not been reproduced yet; a regression test with a BatchNorm in ϕe (checking that the running stats change after a training-mode call) would pin it down. Related: #727.

The other layers with sublayer fields (CGConv, EGNNConv, GATConv, GATv2Conv, GMMConv, GatedGraphConv) have the same return-st pattern, but their constructors only create Dense/Chain of Dense sublayers (plus a GRUCell in GatedGraphConv, whose rng state doesn't appear to be used there), so nothing seems to be lost in practice.

🧔‍♂️🤖 Assisted by AI