Bug Fix: In line 83 of zacharys_karate_club.py, the update function is jax.nn.relu(hk.Linear(5)(n)). Since _ApplyGCN first applies the update function, this corresponds to D^{-0.5}AD^{-0.5}relu(XW) as opposed to relu(D^{-0.5}AD^{-0.5}XW). The update function is now changed to hk.Linear(5) and the relu is applied on the layer output.
Bug Fix: In line 83 of zacharys_karate_club.py, the update function is jax.nn.relu(hk.Linear(5)(n)). Since _ApplyGCN first applies the update function, this corresponds to D^{-0.5}AD^{-0.5}relu(XW) as opposed to relu(D^{-0.5}AD^{-0.5}XW). The update function is now changed to hk.Linear(5) and the relu is applied on the layer output.

PiperOrigin-RevId: 391604165
Change-Id: I88e38d4a0ec02f239357fd53ed78d8e0f27db84b

deepmind/jraph