code: def test_dgi_stateful():
      G = example_graph_random()
      emb_dim = 16
  
      generator = FullBatchNodeGenerator(G)
      corrupted_generator = CorruptedGenerator(generator)
      gen = corrupted_generator.flow(G.nodes())
  
      infomax = DeepGraphInfomax(
          GCN(
              generator=generator,
              activations=["relu"],
              layer_sizes=[emb_dim],
          ),
          corrupted_generator,
      )
  
      model_1 = tf.keras.Model(*infomax.in_out_tensors())
      model_2 = tf.keras.Model(*infomax.in_out_tensors())
  
      # check embeddings are equal before training
      embeddings_1 = tf.keras.Model(*infomax.embedding_model()).predict(
          generator.flow(G.nodes())
      )
      embeddings_2 = tf.keras.Model(*infomax.embedding_model()).predict(
          generator.flow(G.nodes())
      )
  
      np.testing.assert_array_equal(embeddings_1, embeddings_2)
  
      model_1.compile(
          loss=tf.nn.sigmoid_cross_entropy_with_logits, optimizer="Adam"
      )
      model_1.fit(gen)
  
      # check embeddings are still equal after training one model
      embeddings_1 = tf.keras.Model(*infomax.embedding_model()).predict(
          generator.flow(G.nodes())
      )
      embeddings_2 = tf.keras.Model(*infomax.embedding_model()).predict(
          generator.flow(G.nodes())
      )
  
      np.testing.assert_array_equal(embeddings_1, embeddings_2)
  
      model_2.compile(
          loss=tf.nn.sigmoid_cross_entropy_with_logits, optimizer="Adam"
      )
      model_2.fit(gen)
  
      # check embeddings are still equal after training both models
      embeddings_1 = tf.keras.Model(*infomax.embedding_model()).predict(
          generator.flow(G.nodes())
      )
      embeddings_2 = tf.keras.Model(*infomax.embedding_model()).predict(
          generator.flow(G.nodes())
      )
  
      np.testing.assert_array_equal(embeddings_1, embeddings_2)
  
file_path: stellargraph/tests/layer/test_deep_graph_infomax.py
imports: from stellargraph.layer import *
  from stellargraph.mapper import *
  from ..test_utils.graphs import example_graph_random
  from ..None import require_gpu, test_utils
  import tensorflow
  import pytest
  import numpy
label_desc: Use image node_id 3 for calling a global function with example usage: test_dgi_stateful() without return types