code: def _correlation_cost(
      input_a,
      input_b,
      kernel_size,
      max_displacement,
      stride_1,
      stride_2,
      pad,
      data_format="channels_last",
      name=None,
  ):
      """Correlation Cost Volume computation.
  
      See [FlowNet: Learning Optical Flow with Convolutional Networks](https://arxiv.org/abs/1504.06852).
  
      Computes a cost volume using correlation for two inputs. For feature
      maps A, B with spatial dimensions w, h, c it computes
  
        output(a, b) = sum_{l in [-k,k]**2}  < I(a+l), J(b+l) >
  
      where the patches of size K=2d + 1 are centered in position a resp. b.
  
      The output shape is [B, C', H', W'], where
  
        r = max_displacement / stride_2;
        bd = max_displacement + (kernel_size - 1) / 2
        C' = (2 * r + 1) ** 2
        H' = H + 2 * (pad - bd) / stride_1
        W' = W + 2 * (pad - bd) / stride_1
  
      Note: When the data_format requests "channels_last", an additional explicit
        transpose operation is executed.
  
      Args:
        input_a: A `Tensor` of the format specified by `data_format`.
        input_b: A `Tensor` of the format specified by `data_format`.
        kernel_size: An integer specifying the height and width of the
            patch used to compute the per-patch costs.
        max_displacement: An integer specifying the maximum search radius
            for each position.
        stride_1: An integer specifying the stride length in the input.
        stride_2: An integer specifying the stride length in the patch.
        pad: An integer specifying the paddings in height and width.
        data_format: Specifies the data format.
            Possible values are:
            "channels_last" float [batch, height, width, channels]
            "channels_first" float [batch, channels, height, width]
            Defaults to `"channels_last"`.
        name: A name for the operation (optional).
  
      Returns:
        A `Tensor` of the format specified by `data_format`.
      """
  
      with tf.name_scope(name or "correlation_cost"):
          op_call = _correlation_cost_so.ops.addons_correlation_cost
  
          if data_format == "channels_last":
              op_data_format = "NHWC"
          elif data_format == "channels_first":
              op_data_format = "NCHW"
          else:
              raise ValueError(
                  "`data_format` must be either `channels_last` or"
                  "`channels_first`"
              )
  
          ret = op_call(
              input_a,
              input_b,
              kernel_size=kernel_size,
              max_displacement=max_displacement,
              stride_1=stride_1,
              stride_2=stride_2,
              pad=pad,
              data_format=op_data_format,
          )
          if data_format == "channels_last":
              # this is easier to maintain without
              # specializing an additional cuda kernel
              return tf.transpose(ret, [0, 2, 3, 1])
          return ret
  
file_path: addons/tensorflow_addons/layers/optical_flow.py
imports: import tensorflow
  from typeguard import typechecked
  from tensorflow_addons.utils.resource_loader import LazySO
label_desc: Use image node_id 1 for calling a global function with example usage: _correlation_cost(input_a, input_b, kernel_size, max_displacement, stride_1, stride_2, pad, data_format, name) and returns: ret, tf