Repository: tensorflow/addons

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


Discussion:
Hello @facaiy ,

I am new to open source but have good experience in ML. I am interested to contribute to **_addons_** and this seems to be a good starting task.
Could you please assign it to me?

Thanks
Welcome, @PyExtreme :-)
Hi @facaiy , @WindQAQ , I am facing some bazel related issue while running the tests. I am trying to  solve it. I have followed the documentation as well.
It would be great if you could please provide me the commands list in sequence to set up the entire environment and run tests.

Resolving patch:
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 @@ def _correlation_cost_grad(op, grad_output):
 
 @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 @@ def call(self, inputs):
         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,


Test patch: