Upgrade tensorflow and tensorflow_gnn

[jasonrute]

Currently, our code only works with tensorflow < 2.10 and tensorflow_gnn <0.3. Here are the specific issues with upgrading.

tensorflow_gnn==0.3

Causes an issue in prediction with the tf_gnn model. In particular, to reproduce:

• Set pip install tensorflow_gnn==0.3
• Run g2t-server --arch tfgnn ... (no need to interact with Coq)
• Get following error:

  Traceback (most recent call last):
  File "/dccstor/mathzero/debug_memory_leak/venv/bin/g2t-server", line 33, in <module>
  sys.exit(load_entry_point('graph2tac', 'console_scripts', 'g2t-server')())
  File "/dccstor/mathzero/debug_memory_leak/graph2tac/graph2tac/loader/predict_server.py", line 873, in main
  predict = TFGNNPredict(log_dir=Path(args.model).expanduser().absolute(),
  File "/dccstor/mathzero/debug_memory_leak/graph2tac/graph2tac/tfgnn/predict.py", line 180, in __init__
  self.prediction_task = PredictionTask.from_yaml_config(graph_constants=dataset.graph_constants(),
  File "/dccstor/mathzero/debug_memory_leak/graph2tac/graph2tac/tfgnn/tasks.py", line 296, in from_yaml_config
  return prediction_task_constructor(graph_constants=graph_constants, **task_config)
  File "/dccstor/mathzero/debug_memory_leak/graph2tac/graph2tac/tfgnn/tasks.py", line 682, in __init__
  super().__init__(**kwargs)
  File "/dccstor/mathzero/debug_memory_leak/graph2tac/graph2tac/tfgnn/tasks.py", line 483, in __init__
  super().__init__(**kwargs)
  File "/dccstor/mathzero/debug_memory_leak/graph2tac/graph2tac/tfgnn/tasks.py", line 343, in __init__
  super().__init__(**kwargs)
  File "/dccstor/mathzero/debug_memory_leak/graph2tac/graph2tac/tfgnn/tasks.py", line 258, in __init__
  self.graph_embedding = GraphEmbedding(node_label_num=graph_constants.node_label_num,
  File "/dccstor/mathzero/debug_memory_leak/graph2tac/graph2tac/tfgnn/models.py", line 258, in __init__
  self._total_size = tfgnn.keras.layers.TotalSize()
  AttributeError: module 'tensorflow_gnn.keras.layers' has no attribute 'TotalSize'

  in line https://github.com/IBM/graph2tac/blob/1e394a76e805ba866a60a757c774da4715789d5c/graph2tac/tfgnn/models.py#L258

tensorflow==2.10

Causes an issue when training the TF-GNN model. To reproduce:

• pip install tensorflow==2.10
• Train the tfgnn model on multiple gpus
• Get following error:

  Traceback (most recent call last):
  File "/dccstor/mathzero/one_off_tfgnn_runs/2022-09-09_B_tfgnn_stdlib_9-1-split_names-off/scripts/venv/bin/g2t-train-tfgnn", line 8, in <module>
  sys.exit(main())
  File "/dccstor/mathzero/one_off_tfgnn_runs/2022-09-09_B_tfgnn_stdlib_9-1-split_names-off/scripts/venv/lib/python3.10/site-packages/graph2tac/tfgnn/train.py", line 411, in main
  atexit.register(strategy._extended._collective_ops._pool.close)
  AttributeError: 'CollectiveAllReduce' object has no attribute '_pool'

  at line https://github.com/IBM/graph2tac/blob/1e394a76e805ba866a60a757c774da4715789d5c/graph2tac/tfgnn/train.py#L409-L411

[fidel-schaposnik]

The issue with TF-GNN is related to https://github.com/tensorflow/gnn/pull/96 , which removed tfgnn.keras.layers.TotalSize and introduced tfgnn.keras.layers.MakeEmptyFeature. Now it should be possible (and actually easier) to create the empty hidden_state feature for the tfgnn.GraphTensor context, where the node hidden states are later pooled at the end of message-passing.