Add correlation cost layer

WORKSPACE

@@ -1,6 +1,19 @@
+load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
 load("//build_deps/tf_dependency:tf_configure.bzl", "tf_configure")
 load("//build_deps/gpu:cuda_configure.bzl", "cuda_configure")
 
+
+http_archive(
+    name = "cub_archive",
+    build_file = "//build_deps/gpu:cub.BUILD",
+    sha256 = "6bfa06ab52a650ae7ee6963143a0bbc667d6504822cbd9670369b598f18c58c3",
+    strip_prefix = "cub-1.8.0",
+    urls = [
+        "https://storage.googleapis.com/mirror.tensorflow.org/github.com/NVlabs/cub/archive/1.8.0.zip",
+        "https://github.com/NVlabs/cub/archive/1.8.0.zip",
+    ],
+)
+
 tf_configure(
     name = "local_config_tf",
 )

build_deps/gpu/cub.BUILD

@@ -0,0 +1,25 @@
+# Description: CUB library which is a set of primitives for GPU programming.
+
+load("@local_config_cuda//cuda:build_defs.bzl", "cuda_default_copts", "if_cuda")
+
+package(
+    default_visibility = ["//visibility:public"],
+)
+
+licenses(["notice"])  # BSD
+
+filegroup(
+    name = "cub_header_files",
+    srcs = glob([
+        "cub/**",
+    ]),
+)
+
+cc_library(
+    name = "cub",
+    hdrs = if_cuda([":cub_header_files"]),
+    include_prefix = "gpu",
+    deps = [
+        "@local_config_cuda//cuda:cuda_headers",
+    ],
+)

tensorflow_addons/custom_ops/README.md

@@ -6,3 +6,4 @@
 | Image | Ops for image manipulation   |
 | Seq2seq | Ops for seq2seq encoder-decoder framework |
 | Text |  Ops for text processing  |
+| Layers |  Ops for model layers  |

tensorflow_addons/custom_ops/layers/BUILD

@@ -0,0 +1,49 @@
+licenses(["notice"])  # Apache 2.0
+
+package(default_visibility = ["//visibility:public"])
+
+load("@local_config_tf//:build_defs.bzl", "D_GLIBCXX_USE_CXX11_ABI")
+load("@local_config_cuda//cuda:build_defs.bzl", "if_cuda_is_configured", "if_cuda")
+
+cc_binary(
+    name = "_correlation_cost_ops.so",
+    srcs = [
+        "cc/kernels/correlation_cost_op.cc",
+        "cc/kernels/correlation_cost_op.h",
+        "cc/kernels/correlation_cost_op_gpu.cu.cc",
+        "cc/ops/correlation_cost_op.cc",
+    ],
+    copts = [
+        "-pthread",
+        "-std=c++11",
+        D_GLIBCXX_USE_CXX11_ABI,
+    ],
+    linkshared = 1,
+    deps = [
+        "@local_config_tf//:libtensorflow_framework",
+        "@local_config_tf//:tf_header_lib",
+    ] + if_cuda_is_configured([":correlation_cost_ops_gpu"]),
+)
+
+cc_library(
+    name = "correlation_cost_ops_gpu",
+    srcs = [
+        "cc/kernels/correlation_cost_op.h",
+        "cc/kernels/correlation_cost_op_gpu.cu.cc",
+    ],
+    copts = if_cuda_is_configured([
+        "-DGOOGLE_CUDA=1",
+        "-x cuda",
+        "-nvcc_options=relaxed-constexpr",
+        "-nvcc_options=ftz=true",
+    ]),
+    deps = [
+        "@local_config_tf//:libtensorflow_framework",
+        "@local_config_tf//:tf_header_lib",
+    ] + if_cuda_is_configured([
+        "@local_config_cuda//cuda:cuda_libs",
+        "@local_config_cuda//cuda:cuda_headers",
+        "@cub_archive//:cub",
+    ]),
+    alwayslink = 1,
+)