From a26cc18a8b8cfb7d145cb9a26fb30cf7b4cb2868 Mon Sep 17 00:00:00 2001
From: Yanhui Liang <yhliang@google.com>
Date: Thu, 5 Nov 2020 16:46:51 -0800
Subject: [PATCH] Expose `broadcast_weights` as a tf.__internal__ API.

PiperOrigin-RevId: 340952998
Change-Id: If3da39180bc5aa726a93df6cb286eba33c760a4c
---
 tensorflow/python/ops/weights_broadcast_ops.py             | 2 ++
 tensorflow/python/tools/api/generator/api_init_files.bzl   | 1 +
 .../tools/api/golden/v2/tensorflow.__internal__.ops.pbtxt  | 7 +++++++
 .../tools/api/golden/v2/tensorflow.__internal__.pbtxt      | 4 ++++
 4 files changed, 14 insertions(+)
 create mode 100644 tensorflow/tools/api/golden/v2/tensorflow.__internal__.ops.pbtxt

diff --git a/tensorflow/python/ops/weights_broadcast_ops.py b/tensorflow/python/ops/weights_broadcast_ops.py
index 35e93249c31..bfaa7e08fe9 100644
--- a/tensorflow/python/ops/weights_broadcast_ops.py
+++ b/tensorflow/python/ops/weights_broadcast_ops.py
@@ -28,6 +28,7 @@ from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import control_flow_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import sets
+from tensorflow.python.util.tf_export import tf_export
 
 
 def _has_valid_dims(weights_shape, values_shape):
@@ -133,6 +134,7 @@ def assert_broadcastable(weights, values):
     return control_flow_ops.Assert(is_valid_shape, data, name=scope)
 
 
+@tf_export("__internal__.ops.broadcast_weights", v1=[])
 def broadcast_weights(weights, values):
   """Broadcast `weights` to the same shape as `values`.
 
diff --git a/tensorflow/python/tools/api/generator/api_init_files.bzl b/tensorflow/python/tools/api/generator/api_init_files.bzl
index 8dac33cec21..2e0f1ad8536 100644
--- a/tensorflow/python/tools/api/generator/api_init_files.bzl
+++ b/tensorflow/python/tools/api/generator/api_init_files.bzl
@@ -11,6 +11,7 @@ TENSORFLOW_API_INIT_FILES = [
     "__internal__/distribute/combinations/__init__.py",
     "__internal__/distribute/multi_process_runner/__init__.py",
     "__internal__/nest/__init__.py",
+    "__internal__/ops/__init__.py",
     "__internal__/test/__init__.py",
     "__internal__/test/combinations/__init__.py",
     "__internal__/tf2/__init__.py",
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.__internal__.ops.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.__internal__.ops.pbtxt
new file mode 100644
index 00000000000..057ec6f6bdc
--- /dev/null
+++ b/tensorflow/tools/api/golden/v2/tensorflow.__internal__.ops.pbtxt
@@ -0,0 +1,7 @@
+path: "tensorflow.__internal__.ops"
+tf_module {
+  member_method {
+    name: "broadcast_weights"
+    argspec: "args=[\'weights\', \'values\'], varargs=None, keywords=None, defaults=None"
+  }
+}
diff --git a/tensorflow/tools/api/golden/v2/tensorflow.__internal__.pbtxt b/tensorflow/tools/api/golden/v2/tensorflow.__internal__.pbtxt
index 35d23f01a51..6bd4e43b158 100644
--- a/tensorflow/tools/api/golden/v2/tensorflow.__internal__.pbtxt
+++ b/tensorflow/tools/api/golden/v2/tensorflow.__internal__.pbtxt
@@ -20,6 +20,10 @@ tf_module {
     name: "nest"
     mtype: "<type \'module\'>"
   }
+  member {
+    name: "ops"
+    mtype: "<type \'module\'>"
+  }
   member {
     name: "test"
     mtype: "<type \'module\'>"