Export dispatch related function to tf.__internal__ API.

PiperOrigin-RevId: 340368751
Change-Id: I26d3b7b6af2c6f995feba21ea691f0a5128f82ec

tensorflow/python/tools/api/generator/api_init_files.bzl

@@ -6,6 +6,7 @@ TENSORFLOW_API_INIT_FILES = [
     "__init__.py",
     "__internal__/__init__.py",
     "__internal__/decorator/__init__.py",
+    "__internal__/dispatch/__init__.py",
     "__internal__/distribute/__init__.py",
     "__internal__/distribute/combinations/__init__.py",
     "__internal__/distribute/multi_process_runner/__init__.py",

tensorflow/python/util/dispatch.py

@@ -34,6 +34,8 @@ import itertools
 
 from tensorflow.python.util import tf_decorator
 from tensorflow.python.util import tf_inspect
+from tensorflow.python.util.tf_export import tf_export
+
 
 # Private function attribute used to store a list of dispatchers.
 DISPATCH_ATTR = "_tf_dispatchers"
@@ -43,6 +45,7 @@ DISPATCH_ATTR = "_tf_dispatchers"
 _GLOBAL_DISPATCHERS = []
 
 
+@tf_export("__internal__.dispatch.OpDispatcher", v1=[])
 class OpDispatcher(object):
   """Abstract base class for TensorFlow operator dispatchers.
 
@@ -86,6 +89,7 @@ class OpDispatcher(object):
     getattr(op, DISPATCH_ATTR).append(self)
 
 
+@tf_export("__internal__.dispatch.GlobalOpDispatcher", v1=[])
 class GlobalOpDispatcher(object):
   """Abstract base class for TensorFlow global operator dispatchers."""
 
@@ -193,6 +197,7 @@ def add_dispatch_list(target):
   return target
 
 
+@tf_export("__internal__.dispatch.add_dispatch_support", v1=[])
 def add_dispatch_support(target):
   """Decorator that adds a dispatch handling wrapper to an op."""
   def wrapper(*args, **kwargs):

tensorflow/tools/api/golden/v2/tensorflow.__internal__.dispatch.-global-op-dispatcher.pbtxt

@@ -0,0 +1,20 @@
+path: "tensorflow.__internal__.dispatch.GlobalOpDispatcher"
+tf_class {
+  is_instance: "<class \'tensorflow.python.util.dispatch.GlobalOpDispatcher\'>"
+  is_instance: "<type \'object\'>"
+  member {
+    name: "NOT_SUPPORTED"
+    mtype: "<type \'object\'>"
+  }
+  member_method {
+    name: "__init__"
+  }
+  member_method {
+    name: "handle"
+    argspec: "args=[\'self\', \'op\', \'args\', \'kwargs\'], varargs=None, keywords=None, defaults=None"
+  }
+  member_method {
+    name: "register"
+    argspec: "args=[\'self\'], varargs=None, keywords=None, defaults=None"
+  }
+}

tensorflow/tools/api/golden/v2/tensorflow.__internal__.dispatch.-op-dispatcher.pbtxt

@@ -0,0 +1,20 @@
+path: "tensorflow.__internal__.dispatch.OpDispatcher"
+tf_class {
+  is_instance: "<class \'tensorflow.python.util.dispatch.OpDispatcher\'>"
+  is_instance: "<type \'object\'>"
+  member {
+    name: "NOT_SUPPORTED"
+    mtype: "<type \'object\'>"
+  }
+  member_method {
+    name: "__init__"
+  }
+  member_method {
+    name: "handle"
+    argspec: "args=[\'self\', \'args\', \'kwargs\'], varargs=None, keywords=None, defaults=None"
+  }
+  member_method {
+    name: "register"
+    argspec: "args=[\'self\', \'op\'], varargs=None, keywords=None, defaults=None"
+  }
+}

tensorflow/tools/api/golden/v2/tensorflow.__internal__.dispatch.pbtxt

@@ -0,0 +1,15 @@
+path: "tensorflow.__internal__.dispatch"
+tf_module {
+  member {
+    name: "GlobalOpDispatcher"
+    mtype: "<type \'type\'>"
+  }
+  member {
+    name: "OpDispatcher"
+    mtype: "<type \'type\'>"
+  }
+  member_method {
+    name: "add_dispatch_support"
+    argspec: "args=[\'target\'], varargs=None, keywords=None, defaults=None"
+  }
+}

tensorflow/tools/api/golden/v2/tensorflow.__internal__.pbtxt

@@ -8,6 +8,10 @@ tf_module {
     name: "decorator"
     mtype: "<type \'module\'>"
   }
+  member {
+    name: "dispatch"
+    mtype: "<type \'module\'>"
+  }
   member {
     name: "distribute"
     mtype: "<type \'module\'>"