[TF2] Simplify gnn_model_base init and _forward

Currently gnn_model_base and a couple other subclasses have a lot of duplicated code between their `init` functions and their _forward` functions. We should be able to figure out a better way to abstract this out to preserve functionality while making the code simpler.