code: def test_dgi(model_type, sparse):
      base_generator, base_model, nodes = _model_data(
          model_type, sparse
      )
      corrupted_generator = CorruptedGenerator(base_generator)
      gen = corrupted_generator.flow(nodes)
  
      infomax = DeepGraphInfomax(base_model, corrupted_generator)
  
      model = tf.keras.Model(*infomax.in_out_tensors())
      model.compile(
          loss=tf.nn.sigmoid_cross_entropy_with_logits, optimizer="Adam"
      )
      model.fit(gen)
  
      emb_model = tf.keras.Model(*base_model.in_out_tensors())
      embeddings = emb_model.predict(base_generator.flow(nodes))
  
      if isinstance(
          base_generator,
          (FullBatchNodeGenerator, RelationalFullBatchNodeGenerator),
      ):
          assert embeddings.shape == (1, len(nodes), 16)
      else:
          assert embeddings.shape == (len(nodes), 16)
  
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 2 for calling a global function with example usage: test_dgi(model_type, sparse) without return types