code: def test_forward_simple(data_format):
      # We are just testing where the output has vanishing values.
      val_a, val_b = _create_test_data(data_format)
      input_a = tf.constant(val_a, dtype=tf.float32)
      input_b = tf.constant(val_b, dtype=tf.float32)
  
      input_a_tensor = tf.convert_to_tensor(input_a, dtype=tf.float32)
      input_b_tensor = tf.convert_to_tensor(input_b, dtype=tf.float32)
  
      kernel_size = 1
      max_displacement = 2
      stride_1 = 1
      stride_2 = 2
      pad = 4
  
      actual = _forward(
          input_a_tensor,
          input_b_tensor,
          kernel_size=kernel_size,
          max_displacement=max_displacement,
          stride_1=stride_1,
          stride_2=stride_2,
          pad=pad,
          data_format=data_format,
      )
  
      if data_format == "channels_last":
          # NHWC -> NCHW
          actual = tf.transpose(actual, [0, 3, 1, 2])
  
      # We can test fixed ids, as output is independent from data_format
      expected_ids = np.concatenate(
          [
              np.zeros(464),
              np.ones(464),
          ]
      )
      np.testing.assert_allclose(
          tf.where(actual == 0)[:, 0].numpy(), expected_ids
      )
  
      counts = [54, 52, 54, 50, 44, 50, 54, 52, 54]
      expected_ids = np.concatenate(
          [k * np.ones(v) for k, v in enumerate(counts)]
      )
      expected_ids = np.concatenate([expected_ids, expected_ids])
      np.testing.assert_allclose(
          tf.where(actual == 0)[:, 1], expected_ids
      )
      assert actual.shape == (2, 9, 7, 8)
  
file_path: addons/tensorflow_addons/layers/tests/optical_flow_test.py
imports: import pytest
  import numpy
  import tensorflow
  from tensorflow_addons.layers.optical_flow import CorrelationCost
label_desc: Use image node_id 3 for calling a global function with example usage: test_forward_simple(data_format) without return types