Remove framework_headers filegroup, replace use with just collecting

headers from "core:framework" target like we do for framework_headers_lib,
so that it is available in the pip package consistently.

Tested by building the wheel with this change and validating that
the files were in site-packages/tensorflow/include/tensorflow/core/framework,
etc.

Fixes #4996.
Change: 136664737

tensorflow/core/BUILD

@@ -1107,83 +1107,6 @@ cc_header_only_library(
     ],
 )
 
-filegroup(
-    name = "framework_headers",
-    srcs = [
-        "framework/allocator.h",
-        "framework/attr_value_util.h",
-        "framework/bfloat16.h",
-        "framework/cancellation.h",
-        "framework/control_flow.h",
-        "framework/device_base.h",
-        "framework/function.h",
-        "framework/kernel_def_builder.h",
-        "framework/node_def_util.h",
-        "framework/numeric_types.h",
-        "framework/op.h",
-        "framework/op_def_builder.h",
-        "framework/op_def_util.h",
-        "framework/op_kernel.h",
-        "framework/partial_tensor_shape.h",
-        "framework/register_types.h",
-        "framework/rendezvous.h",
-        "framework/selective_registration.h",
-        "framework/session_state.h",
-        "framework/shape_inference.h",
-        "framework/tensor.h",
-        "framework/tensor_reference.h",
-        "framework/tensor_shape.h",
-        "framework/tensor_types.h",
-        "framework/tracking_allocator.h",
-        "framework/type_traits.h",
-        "framework/types.h",
-        "framework/unique_tensor_references.h",
-        "lib/core/errors.h",
-        "lib/core/notification.h",
-        "lib/core/refcount.h",
-        "lib/core/status.h",
-        "lib/core/stringpiece.h",
-        "lib/core/threadpool.h",
-        "lib/gtl/array_slice.h",
-        "lib/gtl/array_slice_internal.h",
-        "lib/gtl/inlined_vector.h",
-        "lib/gtl/manual_constructor.h",
-        "lib/hash/hash.h",
-        "lib/strings/numbers.h",
-        "lib/strings/str_util.h",
-        "lib/strings/strcat.h",
-        "platform/cpu_info.h",
-        "platform/default/dynamic_annotations.h",
-        "platform/default/integral_types.h",
-        "platform/default/logging.h",
-        "platform/default/mutex.h",
-        "platform/default/notification.h",
-        "platform/default/protobuf.h",
-        "platform/default/thread_annotations.h",
-        "platform/dynamic_annotations.h",
-        "platform/env.h",
-        "platform/file_statistics.h",
-        "platform/file_system.h",
-        "platform/fingerprint.h",
-        "platform/logging.h",
-        "platform/macros.h",
-        "platform/mem.h",
-        "platform/mutex.h",
-        "platform/net.h",
-        "platform/notification.h",
-        "platform/platform.h",
-        "platform/prefetch.h",
-        "platform/protobuf.h",
-        "platform/strong_hash.h",
-        "platform/thread_annotations.h",
-        "platform/types.h",
-        "public/session.h",
-        "public/session_options.h",
-        "public/version.h",
-        "util/device_name_utils.h",
-    ],
-)
-
 tf_cuda_library(
     name = "stream_executor",
     srcs = tf_additional_stream_executor_srcs(),

tensorflow/tensorflow.bzl

@@ -630,11 +630,53 @@ def transitive_hdrs(name, deps=[], **kwargs):
   native.filegroup(name=name,
                    srcs=[":" + name + "_gather"])
 
-# Create a header only library that includes all the headers exported by
-# the libraries in deps.
+# The following rules and aspects are used to collect the headers of
+# the direct dependencies of a target.
+def _collect_hdrs_aspect_impl(target, ctx):
+  allhdrs = set()
+  for h in ctx.rule.attr.hdrs:
+    allhdrs = allhdrs | h.files
+  return struct(tf_hdrs=allhdrs)
+
+collect_hdrs_aspect = aspect(
+    implementation=_collect_hdrs_aspect_impl,
+    attr_aspects=["deps"])
+
+def _exposed_hdrs_impl(ctx):
+  all_hdrs = set()
+  for input_dep in ctx.attr.deps:
+    if not hasattr(input_dep, "tf_hdrs"):
+      continue
+    all_hdrs = all_hdrs | input_dep.tf_hdrs
+  return struct(files=all_hdrs)
+
+_exposed_hdrs_rule = rule(
+    _exposed_hdrs_impl,
+    attrs = {
+        "deps": attr.label_list(
+            aspects=[collect_hdrs_aspect],
+            mandatory = True,
+            allow_files = True
+        ),
+    }
+)
+
+def _exposed_hdrs(name, deps=[], **kwargs):
+  _exposed_hdrs_rule(name=name + "_internal", deps=deps)
+  # Defines a filegroup containing all of the headers exposed by
+  # each target in 'deps'.
+  native.filegroup(name=name, srcs=[":" + name + "_internal"])
+
+
+# Create a header only library that includes all the headers exported
+# by the libraries in deps, and exports only the headers of the direct
+# dependencies of entries in 'deps' as a filegroup.
 def cc_header_only_library(name, deps=[], **kwargs):
   _transitive_hdrs(name=name + "_gather",
                    deps=deps)
+
+  _exposed_hdrs(name=name + "_headers", deps=deps)
+
   native.cc_library(name=name,
                     hdrs=[":" + name + "_gather"],
                     **kwargs)

tensorflow/tools/pip_package/BUILD

@@ -5,6 +5,8 @@ package(default_visibility = ["//visibility:private"])
 
 load("//tensorflow:tensorflow.bzl", "transitive_hdrs")
 
+# Bring in the eigen3 and generated protobuf headers that are not part
+# of the direct framework headers lib.
 transitive_hdrs(
     name = "other_headers",
     deps = [
@@ -38,7 +40,7 @@ sh_binary(
         "//tensorflow/contrib/specs:all_files",
         "//tensorflow/contrib/tensor_forest:all_files",
         "//tensorflow/contrib/tensor_forest/hybrid:all_files",
-        "//tensorflow/core:framework_headers",
+        "//tensorflow/core:framework_headers_lib_headers",
         "//tensorflow/examples/tutorials/mnist:package",
         "//tensorflow/models/embedding:package",
         "//tensorflow/models/image/alexnet:all_files",