code: def test_gradients(data_format):
      batch, channels, height, width = 2, 3, 5, 6
      input_a = np.random.randn(batch, channels, height, width).astype(
          np.float32
      )
      input_b = np.random.randn(batch, channels, height, width).astype(
          np.float32
      )
  
      kernel_size = 1
      max_displacement = 2
      stride_1 = 1
      stride_2 = 2
      pad = 4
  
      if data_format == "channels_last":
          input_a = tf.transpose(input_a, [0, 2, 3, 1])
          input_b = tf.transpose(input_b, [0, 2, 3, 1])
  
      input_a_op = tf.convert_to_tensor(input_a)
      input_b_op = tf.convert_to_tensor(input_b)
  
      def correlation_fn(input_a, input_b):
          return CorrelationCost(
              kernel_size=kernel_size,
              max_displacement=max_displacement,
              stride_1=stride_1,
              stride_2=stride_2,
              pad=pad,
              data_format=data_format,
          )([input_a, input_b])
  
      theoretical, numerical = tf.test.compute_gradient(
          correlation_fn, [input_a_op, input_b_op]
      )
  
      np.testing.assert_allclose(
          theoretical[0], numerical[0], atol=1e-3
      )
  
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 4 for calling a global function with example usage: test_gradients(data_format) and returns: CorrelationCost