code: def test_output_shape(mode, num_nodes, num_edges):
      """
      Tests that we can compute the output shape correctly.
  
      :param mode: The data mode to use for this test.
      :param num_nodes: The number of nodes to use for the input.
      :param num_edges: The number of edges to use for the input.
      """
      # Arrange.
      # Create valid-looking input shapes.
      node_features_shape = (num_nodes, F)
      edge_features_shape = (num_edges, S)
      node_laplacian_shape = (num_nodes, num_nodes)
      edge_laplacian_shape = (num_edges, num_edges)
      incidence_shape = (num_nodes, num_edges)
  
      if mode != Modes.SINGLE:
          node_features_shape = (batch_size,) + node_features_shape
          edge_features_shape = (batch_size,) + edge_features_shape
      if mode == Modes.BATCH:
          node_laplacian_shape = (batch_size,) + node_laplacian_shape
          edge_laplacian_shape = (batch_size,) + edge_laplacian_shape
          incidence_shape = (batch_size,) + incidence_shape
  
      input_shape = (
          node_features_shape,
          (node_laplacian_shape, edge_laplacian_shape, incidence_shape),
          edge_features_shape,
      )
  
      # Create the layer.
      layer = CensNetConv(NODE_CHANNELS, EDGE_CHANNELS)
  
      # Act.
      got_output_shape = layer.compute_output_shape(input_shape)
  
      # Assert.
      # The output shape should be the same as the input, but with the correct
      # channel numbers.
      expected_node_feature_shape = node_features_shape[:-1] + (
          NODE_CHANNELS,
      )
      expected_edge_feature_shape = edge_features_shape[:-1] + (
          EDGE_CHANNELS,
      )
      assert got_output_shape == (
          expected_node_feature_shape,
          expected_edge_feature_shape,
      )
  
file_path: spektral/tests/test_layers/convolutional/test_censnet_conv.py
imports: import enum
  import networkx
  import numpy
  import pytest
  from core import A, F, S, batch_size
  from tensorflow.keras import Input, Model
  from spektral.layers import CensNetConv
label_desc: Use image node_id 3 for calling a global function with example usage: test_output_shape(mode, num_nodes, num_edges) without return types