We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Add docstring for Correlation cost layer
1. Hide [correlation_cost](https://github.com/tensorflow/addons/blob/cffee80779e606870ebf2b7fce241eeb870053f5/tensorflow_addons/layers/optical_flow.py#L30) function, and make it private.
2. Add document for [CorrelationCost](https://github.com/tensorflow/addons/blob/cffee80779e606870ebf2b7fce241eeb870053f5/tensorflow_addons/layers/optical_flow.py#L143)

For details please check the discussion in https://github.com/tensorflow/addons/pull/207

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `tensorflow_addons/layers/optical_flow.py`
Content:
```
1 # Copyright 2019 The TensorFlow Authors. All Rights Reserved.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #    http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 # ============================================================================
15 """Tensorflow op performing correlation cost operation."""
16 
17 from __future__ import absolute_import
18 from __future__ import division
19 from __future__ import print_function
20 
21 import tensorflow as tf
22 from tensorflow_addons.utils import keras_utils
23 from tensorflow_addons.utils.resource_loader import get_path_to_datafile
24 
25 _correlation_cost_op_so = tf.load_op_library(
26     get_path_to_datafile("custom_ops/layers/_correlation_cost_ops.so"))
27 
28 
29 @tf.function
30 def correlation_cost(input_a,
31                      input_b,
32                      kernel_size,
33                      max_displacement,
34                      stride_1,
35                      stride_2,
36                      pad,
37                      data_format='channels_last',
38                      name=None):
39     """Correlation Cost Volume computation.
40 
41     "FlowNet: Learning Optical Flow with Convolutional Networks"
42     Philipp Fischer, Alexey Dosovitskiy, Eddy Ilg, Philip Hausser,
43     Caner Hazirbas, Vladimir Golkov, Patrick van der Smagt,
44     Daniel Cremers, Thomas Brox. https://arxiv.org/abs/1504.06852
45 
46     Computes a cost volume using correlation for two inputs. For feature
47     maps A, B with spatial dimensions w, h, c it computes
48 
49       output(a, b) = sum_{l in [-k,k]**2}  < I(a+l), J(b+l) >
50 
51     where the patches of size K=2d + 1 are centered in position a resp. b.
52 
53     The output shape is [B, C', H', W'], where
54 
55       r = max_displacement / stride_2;
56       bd = max_displacement + (kernel_size - 1) / 2
57       C' = (2 * r + 1) ** 2
58       H' = H + 2 * (pad - bd) / stride_1
59       W' = W + 2 * (pad - bd) / stride_1
60 
61     Note: When the data_format requests "channels_last", an additional explicit
62       transpose operation is executed.
63 
64     Args:
65       input_a: A `Tensor` of the format specified by `data_format`.
66       input_b: A `Tensor` of the format specified by `data_format`.
67       kernel_size: An integer specifying the height and width of the
68           patch used to compute the per-patch costs.
69       max_displacement: An integer specifying the maximum search radius
70           for each position.
71       stride_1: An integer specifying the stride length in the input.
72       stride_2: An integer specifying the stride length in the patch.
73       pad: An integer specifying the paddings in height and width.
74       data_format: Specifies the data format.
75           Possible values are:
76           "channels_last" float [batch, height, width, channels]
77           "channels_first" float [batch, channels, height, width]
78           Defaults to `"channels_last"`.
79       name: A name for the operation (optional).
80 
81     Returns:
82       A `Tensor` of the format specified by `data_format`.
83     """
84 
85     with tf.name_scope(name or "correlation_cost"):
86         op_call = _correlation_cost_op_so.addons_correlation_cost
87 
88         if data_format == "channels_last":
89             op_data_format = "NHWC"
90         elif data_format == "channels_first":
91             op_data_format = "NCHW"
92         else:
93             raise ValueError("`data_format` must be either `channels_last` or"
94                              "`channels_first`")
95 
96         ret = op_call(
97             input_a,
98             input_b,
99             kernel_size=kernel_size,
100             max_displacement=max_displacement,
101             stride_1=stride_1,
102             stride_2=stride_2,
103             pad=pad,
104             data_format=op_data_format)
105         if data_format == 'channels_last':
106             # this is easier to maintain without
107             # specializing an additional cuda kernel
108             return tf.transpose(ret, [0, 2, 3, 1])
109         return ret
110 
111 
112 @tf.RegisterGradient("Addons>CorrelationCost")
113 def _correlation_cost_grad(op, grad_output):
114     kernel_size = op.get_attr("kernel_size")
115     max_displacement = op.get_attr("max_displacement")
116     stride_1 = op.get_attr("stride_1")
117     stride_2 = op.get_attr("stride_2")
118     pad = op.get_attr("pad")
119     data_format = op.get_attr("data_format")
120 
121     input_a = tf.convert_to_tensor(op.inputs[0], name="input_a")
122     input_b = tf.convert_to_tensor(op.inputs[1], name="input_b")
123     grad_output_tensor = tf.convert_to_tensor(grad_output, name="grad_output")
124 
125     op_call = _correlation_cost_op_so.addons_correlation_cost_grad
126     grads = op_call(
127         input_a,
128         input_b,
129         grad_output_tensor,
130         kernel_size=kernel_size,
131         max_displacement=max_displacement,
132         stride_1=stride_1,
133         stride_2=stride_2,
134         pad=pad,
135         data_format=data_format)
136 
137     grad_input_a = tf.convert_to_tensor(grads[0], name="grad_input_a")
138     grad_input_b = tf.convert_to_tensor(grads[1], name="grad_input_b")
139     return [grad_input_a, grad_input_b]
140 
141 
142 @keras_utils.register_keras_custom_object
143 class CorrelationCost(tf.keras.layers.Layer):
144     def __init__(self, kernel_size, max_displacement, stride_1, stride_2, pad,
145                  data_format, **kwargs):
146         self.kernel_size = kernel_size
147         self.max_displacement = max_displacement
148         self.stride_1 = stride_1
149         self.stride_2 = stride_2
150         self.pad = pad
151 
152         if data_format != "channels_last" and data_format != "channels_first":
153             raise ValueError("`data_format` must be either `channels_last` or"
154                              "`channels_first`, instead got %s" % data_format)
155 
156         self.data_format = data_format
157 
158         super(CorrelationCost, self).__init__(**kwargs)
159 
160     def build(self, input_shape):
161         if not isinstance(input_shape, list):
162             raise ValueError("Input must be a list of two Tensors to process")
163         super(CorrelationCost, self).build(input_shape)
164 
165     def call(self, inputs):
166         if not isinstance(inputs, list):
167             raise ValueError("Input must be a list of two Tensors to process")
168 
169         input_a = tf.convert_to_tensor(inputs[0])
170         input_b = tf.convert_to_tensor(inputs[1])
171 
172         return correlation_cost(
173             input_a,
174             input_b,
175             kernel_size=self.kernel_size,
176             max_displacement=self.max_displacement,
177             stride_1=self.stride_1,
178             stride_2=self.stride_2,
179             pad=self.pad,
180             data_format=self.data_format)
181 
182     def compute_output_shape(self, input_shape):
183         assert isinstance(input_shape, list)
184         n = input_shape[0][0]
185         r = self.max_displacement / self.stride_2
186         bd = self.max_displacement + (self.kernel_size - 1) / 2
187         output_c = (2 * r + 1)**2
188 
189         if self.data_format == "channels_first":
190             output_h = input_shape[0][1] + 2 * (self.pad - bd) / self.stride_1
191             output_w = input_shape[0][2] + 2 * (self.pad - bd) / self.stride_1
192             return [n, output_c, output_h, output_w]
193 
194         elif self.data_format == "channels_last":
195             output_h = input_shape[0][0] + 2 * (self.pad - bd) / self.stride_1
196             output_w = input_shape[0][1] + 2 * (self.pad - bd) / self.stride_1
197             return [n, output_h, output_w, output_c]
198         else:
199             raise ValueError("`data_format` must be either `channels_last` or"
200                              "`channels_first`")
201 
202     def get_config(self):
203         config = {
204             'kernel_size': self.kernel_size,
205             'max_displacement': self.max_displacement,
206             'stride_1': self.stride_1,
207             'stride_2': self.stride_2,
208             'pad': self.pad,
209             'data_format': self.data_format
210         }
211 
212         base_config = super(CorrelationCost, self).get_config()
213         return dict(list(base_config.items()) + list(config.items()))
214
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/tensorflow_addons/layers/optical_flow.py b/tensorflow_addons/layers/optical_flow.py
--- a/tensorflow_addons/layers/optical_flow.py
+++ b/tensorflow_addons/layers/optical_flow.py
@@ -27,15 +27,15 @@
 
 
 @tf.function
-def correlation_cost(input_a,
-                     input_b,
-                     kernel_size,
-                     max_displacement,
-                     stride_1,
-                     stride_2,
-                     pad,
-                     data_format='channels_last',
-                     name=None):
+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.
 
     "FlowNet: Learning Optical Flow with Convolutional Networks"
@@ -141,6 +141,27 @@
 
 @keras_utils.register_keras_custom_object
 class CorrelationCost(tf.keras.layers.Layer):
+    """Correlation Cost Layer.
+
+    This layer implements the correlation operation from FlowNet Learning
+    Optical Flow with Convolutional Networks (Fischer et al.):
+    https://arxiv.org/abs/1504.06
+
+    Args:
+        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"`.
+    """
+
     def __init__(self, kernel_size, max_displacement, stride_1, stride_2, pad,
                  data_format, **kwargs):
         self.kernel_size = kernel_size
@@ -169,7 +190,7 @@
         input_a = tf.convert_to_tensor(inputs[0])
         input_b = tf.convert_to_tensor(inputs[1])
 
-        return correlation_cost(
+        return _correlation_cost(
             input_a,
             input_b,
             kernel_size=self.kernel_size,

{"golden_diff": "diff --git a/tensorflow_addons/layers/optical_flow.py b/tensorflow_addons/layers/optical_flow.py\n--- a/tensorflow_addons/layers/optical_flow.py\n+++ b/tensorflow_addons/layers/optical_flow.py\n@@ -27,15 +27,15 @@\n \n \n @tf.function\n-def correlation_cost(input_a,\n-                     input_b,\n-                     kernel_size,\n-                     max_displacement,\n-                     stride_1,\n-                     stride_2,\n-                     pad,\n-                     data_format='channels_last',\n-                     name=None):\n+def _correlation_cost(input_a,\n+                      input_b,\n+                      kernel_size,\n+                      max_displacement,\n+                      stride_1,\n+                      stride_2,\n+                      pad,\n+                      data_format='channels_last',\n+                      name=None):\n     \"\"\"Correlation Cost Volume computation.\n \n     \"FlowNet: Learning Optical Flow with Convolutional Networks\"\n@@ -141,6 +141,27 @@\n \n @keras_utils.register_keras_custom_object\n class CorrelationCost(tf.keras.layers.Layer):\n+    \"\"\"Correlation Cost Layer.\n+\n+    This layer implements the correlation operation from FlowNet Learning\n+    Optical Flow with Convolutional Networks (Fischer et al.):\n+    https://arxiv.org/abs/1504.06\n+\n+    Args:\n+        kernel_size: An integer specifying the height and width of the\n+            patch used to compute the per-patch costs.\n+        max_displacement: An integer specifying the maximum search radius\n+            for each position.\n+        stride_1: An integer specifying the stride length in the input.\n+        stride_2: An integer specifying the stride length in the patch.\n+        pad: An integer specifying the paddings in height and width.\n+        data_format: Specifies the data format.\n+            Possible values are:\n+                \"channels_last\" float [batch, height, width, channels]\n+                \"channels_first\" float [batch, channels, height, width]\n+                Defaults to `\"channels_last\"`.\n+    \"\"\"\n+\n     def __init__(self, kernel_size, max_displacement, stride_1, stride_2, pad,\n                  data_format, **kwargs):\n         self.kernel_size = kernel_size\n@@ -169,7 +190,7 @@\n         input_a = tf.convert_to_tensor(inputs[0])\n         input_b = tf.convert_to_tensor(inputs[1])\n \n-        return correlation_cost(\n+        return _correlation_cost(\n             input_a,\n             input_b,\n             kernel_size=self.kernel_size,\n", "issue": "Add docstring for Correlation cost layer\n1. Hide [correlation_cost](https://github.com/tensorflow/addons/blob/cffee80779e606870ebf2b7fce241eeb870053f5/tensorflow_addons/layers/optical_flow.py#L30) function, and make it private.\r\n2. Add document for [CorrelationCost](https://github.com/tensorflow/addons/blob/cffee80779e606870ebf2b7fce241eeb870053f5/tensorflow_addons/layers/optical_flow.py#L143)\r\n\r\nFor details please check the discussion in https://github.com/tensorflow/addons/pull/207\n", "before_files": [{"content": "# Copyright 2019 The TensorFlow Authors. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#    http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ============================================================================\n\"\"\"Tensorflow op performing correlation cost operation.\"\"\"\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport tensorflow as tf\nfrom tensorflow_addons.utils import keras_utils\nfrom tensorflow_addons.utils.resource_loader import get_path_to_datafile\n\n_correlation_cost_op_so = tf.load_op_library(\n    get_path_to_datafile(\"custom_ops/layers/_correlation_cost_ops.so\"))\n\n\n@tf.function\ndef correlation_cost(input_a,\n                     input_b,\n                     kernel_size,\n                     max_displacement,\n                     stride_1,\n                     stride_2,\n                     pad,\n                     data_format='channels_last',\n                     name=None):\n    \"\"\"Correlation Cost Volume computation.\n\n    \"FlowNet: Learning Optical Flow with Convolutional Networks\"\n    Philipp Fischer, Alexey Dosovitskiy, Eddy Ilg, Philip Hausser,\n    Caner Hazirbas, Vladimir Golkov, Patrick van der Smagt,\n    Daniel Cremers, Thomas Brox. https://arxiv.org/abs/1504.06852\n\n    Computes a cost volume using correlation for two inputs. For feature\n    maps A, B with spatial dimensions w, h, c it computes\n\n      output(a, b) = sum_{l in [-k,k]**2}  < I(a+l), J(b+l) >\n\n    where the patches of size K=2d + 1 are centered in position a resp. b.\n\n    The output shape is [B, C', H', W'], where\n\n      r = max_displacement / stride_2;\n      bd = max_displacement + (kernel_size - 1) / 2\n      C' = (2 * r + 1) ** 2\n      H' = H + 2 * (pad - bd) / stride_1\n      W' = W + 2 * (pad - bd) / stride_1\n\n    Note: When the data_format requests \"channels_last\", an additional explicit\n      transpose operation is executed.\n\n    Args:\n      input_a: A `Tensor` of the format specified by `data_format`.\n      input_b: A `Tensor` of the format specified by `data_format`.\n      kernel_size: An integer specifying the height and width of the\n          patch used to compute the per-patch costs.\n      max_displacement: An integer specifying the maximum search radius\n          for each position.\n      stride_1: An integer specifying the stride length in the input.\n      stride_2: An integer specifying the stride length in the patch.\n      pad: An integer specifying the paddings in height and width.\n      data_format: Specifies the data format.\n          Possible values are:\n          \"channels_last\" float [batch, height, width, channels]\n          \"channels_first\" float [batch, channels, height, width]\n          Defaults to `\"channels_last\"`.\n      name: A name for the operation (optional).\n\n    Returns:\n      A `Tensor` of the format specified by `data_format`.\n    \"\"\"\n\n    with tf.name_scope(name or \"correlation_cost\"):\n        op_call = _correlation_cost_op_so.addons_correlation_cost\n\n        if data_format == \"channels_last\":\n            op_data_format = \"NHWC\"\n        elif data_format == \"channels_first\":\n            op_data_format = \"NCHW\"\n        else:\n            raise ValueError(\"`data_format` must be either `channels_last` or\"\n                             \"`channels_first`\")\n\n        ret = op_call(\n            input_a,\n            input_b,\n            kernel_size=kernel_size,\n            max_displacement=max_displacement,\n            stride_1=stride_1,\n            stride_2=stride_2,\n            pad=pad,\n            data_format=op_data_format)\n        if data_format == 'channels_last':\n            # this is easier to maintain without\n            # specializing an additional cuda kernel\n            return tf.transpose(ret, [0, 2, 3, 1])\n        return ret\n\n\n@tf.RegisterGradient(\"Addons>CorrelationCost\")\ndef _correlation_cost_grad(op, grad_output):\n    kernel_size = op.get_attr(\"kernel_size\")\n    max_displacement = op.get_attr(\"max_displacement\")\n    stride_1 = op.get_attr(\"stride_1\")\n    stride_2 = op.get_attr(\"stride_2\")\n    pad = op.get_attr(\"pad\")\n    data_format = op.get_attr(\"data_format\")\n\n    input_a = tf.convert_to_tensor(op.inputs[0], name=\"input_a\")\n    input_b = tf.convert_to_tensor(op.inputs[1], name=\"input_b\")\n    grad_output_tensor = tf.convert_to_tensor(grad_output, name=\"grad_output\")\n\n    op_call = _correlation_cost_op_so.addons_correlation_cost_grad\n    grads = op_call(\n        input_a,\n        input_b,\n        grad_output_tensor,\n        kernel_size=kernel_size,\n        max_displacement=max_displacement,\n        stride_1=stride_1,\n        stride_2=stride_2,\n        pad=pad,\n        data_format=data_format)\n\n    grad_input_a = tf.convert_to_tensor(grads[0], name=\"grad_input_a\")\n    grad_input_b = tf.convert_to_tensor(grads[1], name=\"grad_input_b\")\n    return [grad_input_a, grad_input_b]\n\n\n@keras_utils.register_keras_custom_object\nclass CorrelationCost(tf.keras.layers.Layer):\n    def __init__(self, kernel_size, max_displacement, stride_1, stride_2, pad,\n                 data_format, **kwargs):\n        self.kernel_size = kernel_size\n        self.max_displacement = max_displacement\n        self.stride_1 = stride_1\n        self.stride_2 = stride_2\n        self.pad = pad\n\n        if data_format != \"channels_last\" and data_format != \"channels_first\":\n            raise ValueError(\"`data_format` must be either `channels_last` or\"\n                             \"`channels_first`, instead got %s\" % data_format)\n\n        self.data_format = data_format\n\n        super(CorrelationCost, self).__init__(**kwargs)\n\n    def build(self, input_shape):\n        if not isinstance(input_shape, list):\n            raise ValueError(\"Input must be a list of two Tensors to process\")\n        super(CorrelationCost, self).build(input_shape)\n\n    def call(self, inputs):\n        if not isinstance(inputs, list):\n            raise ValueError(\"Input must be a list of two Tensors to process\")\n\n        input_a = tf.convert_to_tensor(inputs[0])\n        input_b = tf.convert_to_tensor(inputs[1])\n\n        return correlation_cost(\n            input_a,\n            input_b,\n            kernel_size=self.kernel_size,\n            max_displacement=self.max_displacement,\n            stride_1=self.stride_1,\n            stride_2=self.stride_2,\n            pad=self.pad,\n            data_format=self.data_format)\n\n    def compute_output_shape(self, input_shape):\n        assert isinstance(input_shape, list)\n        n = input_shape[0][0]\n        r = self.max_displacement / self.stride_2\n        bd = self.max_displacement + (self.kernel_size - 1) / 2\n        output_c = (2 * r + 1)**2\n\n        if self.data_format == \"channels_first\":\n            output_h = input_shape[0][1] + 2 * (self.pad - bd) / self.stride_1\n            output_w = input_shape[0][2] + 2 * (self.pad - bd) / self.stride_1\n            return [n, output_c, output_h, output_w]\n\n        elif self.data_format == \"channels_last\":\n            output_h = input_shape[0][0] + 2 * (self.pad - bd) / self.stride_1\n            output_w = input_shape[0][1] + 2 * (self.pad - bd) / self.stride_1\n            return [n, output_h, output_w, output_c]\n        else:\n            raise ValueError(\"`data_format` must be either `channels_last` or\"\n                             \"`channels_first`\")\n\n    def get_config(self):\n        config = {\n            'kernel_size': self.kernel_size,\n            'max_displacement': self.max_displacement,\n            'stride_1': self.stride_1,\n            'stride_2': self.stride_2,\n            'pad': self.pad,\n            'data_format': self.data_format\n        }\n\n        base_config = super(CorrelationCost, self).get_config()\n        return dict(list(base_config.items()) + list(config.items()))\n", "path": "tensorflow_addons/layers/optical_flow.py"}], "after_files": [{"content": "# Copyright 2019 The TensorFlow Authors. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#    http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ============================================================================\n\"\"\"Tensorflow op performing correlation cost operation.\"\"\"\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport tensorflow as tf\nfrom tensorflow_addons.utils import keras_utils\nfrom tensorflow_addons.utils.resource_loader import get_path_to_datafile\n\n_correlation_cost_op_so = tf.load_op_library(\n    get_path_to_datafile(\"custom_ops/layers/_correlation_cost_ops.so\"))\n\n\n@tf.function\ndef _correlation_cost(input_a,\n                      input_b,\n                      kernel_size,\n                      max_displacement,\n                      stride_1,\n                      stride_2,\n                      pad,\n                      data_format='channels_last',\n                      name=None):\n    \"\"\"Correlation Cost Volume computation.\n\n    \"FlowNet: Learning Optical Flow with Convolutional Networks\"\n    Philipp Fischer, Alexey Dosovitskiy, Eddy Ilg, Philip Hausser,\n    Caner Hazirbas, Vladimir Golkov, Patrick van der Smagt,\n    Daniel Cremers, Thomas Brox. https://arxiv.org/abs/1504.06852\n\n    Computes a cost volume using correlation for two inputs. For feature\n    maps A, B with spatial dimensions w, h, c it computes\n\n      output(a, b) = sum_{l in [-k,k]**2}  < I(a+l), J(b+l) >\n\n    where the patches of size K=2d + 1 are centered in position a resp. b.\n\n    The output shape is [B, C', H', W'], where\n\n      r = max_displacement / stride_2;\n      bd = max_displacement + (kernel_size - 1) / 2\n      C' = (2 * r + 1) ** 2\n      H' = H + 2 * (pad - bd) / stride_1\n      W' = W + 2 * (pad - bd) / stride_1\n\n    Note: When the data_format requests \"channels_last\", an additional explicit\n      transpose operation is executed.\n\n    Args:\n      input_a: A `Tensor` of the format specified by `data_format`.\n      input_b: A `Tensor` of the format specified by `data_format`.\n      kernel_size: An integer specifying the height and width of the\n          patch used to compute the per-patch costs.\n      max_displacement: An integer specifying the maximum search radius\n          for each position.\n      stride_1: An integer specifying the stride length in the input.\n      stride_2: An integer specifying the stride length in the patch.\n      pad: An integer specifying the paddings in height and width.\n      data_format: Specifies the data format.\n          Possible values are:\n          \"channels_last\" float [batch, height, width, channels]\n          \"channels_first\" float [batch, channels, height, width]\n          Defaults to `\"channels_last\"`.\n      name: A name for the operation (optional).\n\n    Returns:\n      A `Tensor` of the format specified by `data_format`.\n    \"\"\"\n\n    with tf.name_scope(name or \"correlation_cost\"):\n        op_call = _correlation_cost_op_so.addons_correlation_cost\n\n        if data_format == \"channels_last\":\n            op_data_format = \"NHWC\"\n        elif data_format == \"channels_first\":\n            op_data_format = \"NCHW\"\n        else:\n            raise ValueError(\"`data_format` must be either `channels_last` or\"\n                             \"`channels_first`\")\n\n        ret = op_call(\n            input_a,\n            input_b,\n            kernel_size=kernel_size,\n            max_displacement=max_displacement,\n            stride_1=stride_1,\n            stride_2=stride_2,\n            pad=pad,\n            data_format=op_data_format)\n        if data_format == 'channels_last':\n            # this is easier to maintain without\n            # specializing an additional cuda kernel\n            return tf.transpose(ret, [0, 2, 3, 1])\n        return ret\n\n\n@tf.RegisterGradient(\"Addons>CorrelationCost\")\ndef _correlation_cost_grad(op, grad_output):\n    kernel_size = op.get_attr(\"kernel_size\")\n    max_displacement = op.get_attr(\"max_displacement\")\n    stride_1 = op.get_attr(\"stride_1\")\n    stride_2 = op.get_attr(\"stride_2\")\n    pad = op.get_attr(\"pad\")\n    data_format = op.get_attr(\"data_format\")\n\n    input_a = tf.convert_to_tensor(op.inputs[0], name=\"input_a\")\n    input_b = tf.convert_to_tensor(op.inputs[1], name=\"input_b\")\n    grad_output_tensor = tf.convert_to_tensor(grad_output, name=\"grad_output\")\n\n    op_call = _correlation_cost_op_so.addons_correlation_cost_grad\n    grads = op_call(\n        input_a,\n        input_b,\n        grad_output_tensor,\n        kernel_size=kernel_size,\n        max_displacement=max_displacement,\n        stride_1=stride_1,\n        stride_2=stride_2,\n        pad=pad,\n        data_format=data_format)\n\n    grad_input_a = tf.convert_to_tensor(grads[0], name=\"grad_input_a\")\n    grad_input_b = tf.convert_to_tensor(grads[1], name=\"grad_input_b\")\n    return [grad_input_a, grad_input_b]\n\n\n@keras_utils.register_keras_custom_object\nclass CorrelationCost(tf.keras.layers.Layer):\n    \"\"\"Correlation Cost Layer.\n\n    This layer implements the correlation operation from FlowNet Learning\n    Optical Flow with Convolutional Networks (Fischer et al.):\n    https://arxiv.org/abs/1504.06\n\n    Args:\n        kernel_size: An integer specifying the height and width of the\n            patch used to compute the per-patch costs.\n        max_displacement: An integer specifying the maximum search radius\n            for each position.\n        stride_1: An integer specifying the stride length in the input.\n        stride_2: An integer specifying the stride length in the patch.\n        pad: An integer specifying the paddings in height and width.\n        data_format: Specifies the data format.\n            Possible values are:\n                \"channels_last\" float [batch, height, width, channels]\n                \"channels_first\" float [batch, channels, height, width]\n                Defaults to `\"channels_last\"`.\n    \"\"\"\n\n    def __init__(self, kernel_size, max_displacement, stride_1, stride_2, pad,\n                 data_format, **kwargs):\n        self.kernel_size = kernel_size\n        self.max_displacement = max_displacement\n        self.stride_1 = stride_1\n        self.stride_2 = stride_2\n        self.pad = pad\n\n        if data_format != \"channels_last\" and data_format != \"channels_first\":\n            raise ValueError(\"`data_format` must be either `channels_last` or\"\n                             \"`channels_first`, instead got %s\" % data_format)\n\n        self.data_format = data_format\n\n        super(CorrelationCost, self).__init__(**kwargs)\n\n    def build(self, input_shape):\n        if not isinstance(input_shape, list):\n            raise ValueError(\"Input must be a list of two Tensors to process\")\n        super(CorrelationCost, self).build(input_shape)\n\n    def call(self, inputs):\n        if not isinstance(inputs, list):\n            raise ValueError(\"Input must be a list of two Tensors to process\")\n\n        input_a = tf.convert_to_tensor(inputs[0])\n        input_b = tf.convert_to_tensor(inputs[1])\n\n        return _correlation_cost(\n            input_a,\n            input_b,\n            kernel_size=self.kernel_size,\n            max_displacement=self.max_displacement,\n            stride_1=self.stride_1,\n            stride_2=self.stride_2,\n            pad=self.pad,\n            data_format=self.data_format)\n\n    def compute_output_shape(self, input_shape):\n        assert isinstance(input_shape, list)\n        n = input_shape[0][0]\n        r = self.max_displacement / self.stride_2\n        bd = self.max_displacement + (self.kernel_size - 1) / 2\n        output_c = (2 * r + 1)**2\n\n        if self.data_format == \"channels_first\":\n            output_h = input_shape[0][1] + 2 * (self.pad - bd) / self.stride_1\n            output_w = input_shape[0][2] + 2 * (self.pad - bd) / self.stride_1\n            return [n, output_c, output_h, output_w]\n\n        elif self.data_format == \"channels_last\":\n            output_h = input_shape[0][0] + 2 * (self.pad - bd) / self.stride_1\n            output_w = input_shape[0][1] + 2 * (self.pad - bd) / self.stride_1\n            return [n, output_h, output_w, output_c]\n        else:\n            raise ValueError(\"`data_format` must be either `channels_last` or\"\n                             \"`channels_first`\")\n\n    def get_config(self):\n        config = {\n            'kernel_size': self.kernel_size,\n            'max_displacement': self.max_displacement,\n            'stride_1': self.stride_1,\n            'stride_2': self.stride_2,\n            'pad': self.pad,\n            'data_format': self.data_format\n        }\n\n        base_config = super(CorrelationCost, self).get_config()\n        return dict(list(base_config.items()) + list(config.items()))\n", "path": "tensorflow_addons/layers/optical_flow.py"}]}