code: def test_distmult(knowledge_graph):
      # this test creates a random untrained model and predicts every possible edge in the graph, and
      # compares that to a direct implementation of the scoring method in the paper
      gen = KGTripleGenerator(knowledge_graph, 3)
  
      # use a random initializer with a large range, so that any differences are obvious
      init = initializers.RandomUniform(-1, 1)
      distmult_model = DistMult(gen, 5, embeddings_initializer=init)
      x_inp, x_out = distmult_model.in_out_tensors()
  
      model = Model(x_inp, x_out)
  
      model.compile(loss=tf_losses.BinaryCrossentropy(from_logits=True))
  
      every_edge = itertools.product(
          knowledge_graph.nodes(),
          knowledge_graph._edges.types.pandas_index,
          knowledge_graph.nodes(),
      )
      df = triple_df(*every_edge)
  
      # check the model can be trained on a few (uneven) batches
      model.fit(
          gen.flow(df.iloc[:7], negative_samples=2),
          validation_data=gen.flow(df.iloc[7:14], negative_samples=3),
      )
  
      # compute the exact values based on the model by extracting the embeddings for each element and
      # doing the y_(e_1)^T M_r y_(e_2) = <e_1, w_r, e_2> inner product
      s_idx = knowledge_graph.node_ids_to_ilocs(df.source)
      r_idx = knowledge_graph._edges.types.to_iloc(df.label)
      o_idx = knowledge_graph.node_ids_to_ilocs(df.target)
  
      nodes, edge_types = distmult_model.embeddings()
      # the rows correspond to the embeddings for the given edge, so we can do bulk operations
      e_s = nodes[s_idx, :]
      w_r = edge_types[r_idx, :]
      e_o = nodes[o_idx, :]
      actual = (e_s * w_r * e_o).sum(axis=1)
  
      # predict every edge using the model
      prediction = model.predict(gen.flow(df))
  
      # (use an absolute tolerance to allow for catastrophic cancellation around very small values)
      np.testing.assert_allclose(
          prediction[:, 0], actual, rtol=1e-3, atol=1e-14
      )
  
      # the model is stateful (i.e. it holds the weights permanently) so the predictions with a second
      # 'build' should be the same as the original one
      model2 = Model(*distmult_model.in_out_tensors())
      prediction2 = model2.predict(gen.flow(df))
      np.testing.assert_array_equal(prediction, prediction2)
  
file_path: stellargraph/tests/layer/test_knowledge_graph.py
imports: import itertools
  import pytest
  import pandas
  import numpy
  import tensorflow
  from tensorflow.keras import Model, initializers, losses, layers
  from stellargraph import StellarGraph, StellarDiGraph, losses
  from stellargraph.mapper.knowledge_graph import KGTripleGenerator
  from stellargraph.layer.knowledge_graph import KGModel, KGScore, ComplEx, DistMult, RotatE, RotE, RotH, _ranks_from_score_columns
  from ..None import test_utils
  from ..test_utils.graphs import knowledge_graph
label_desc: Use image node_id 3 for calling a global function with example usage: test_distmult(knowledge_graph) without return types