I'm getting this error running the demo. Google suggests this was removed in Sonnet 2 (TF2), and replaced by Module (see comment half way down here google-deepmind/sonnet#128). Full error:
2021-02-25 13:58:31.737640: I tensorflow/stream_executor/platform/default/dso_loader.cc:49] Successfully opened dynamic library cudart64_110.dll
Traceback (most recent call last):
File "models.py", line 44, in
class MLPGraphIndependent(snt.AbstractModule):
AttributeError: module 'sonnet' has no attribute 'AbstractModule'