chore: import upstream snapshot with attribution

tensorflow/python/tools/api/generator2/BUILD

@@ -0,0 +1,45 @@
+load("@bazel_skylib//:bzl_library.bzl", "bzl_library")
+load("//tensorflow:pytype.default.bzl", "pytype_strict_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//tensorflow/python/tools/api/generator2:__subpackages__"],
+    licenses = ["notice"],
+)
+
+config_setting(
+    name = "static_gen",
+    define_values = {
+        "TF_API_INIT_LOADING": "static",
+    },
+)
+
+bzl_library(
+    name = "apis_bzl",
+    srcs = ["apis.bzl"],
+    visibility = ["//visibility:private"],
+    deps = [":patterns_bzl"],
+)
+
+bzl_library(
+    name = "generate_api_bzl",
+    srcs = ["generate_api.bzl"],
+    deps = [
+        ":apis_bzl",
+        ":patterns_bzl",
+        "//tensorflow/python/tools/api/generator:api_init_files",
+        "@bazel_skylib//lib:paths",
+    ],
+)
+
+bzl_library(
+    name = "patterns_bzl",
+    srcs = ["patterns.bzl"],
+    visibility = ["//visibility:private"],
+)
+
+pytype_strict_library(
+    name = "docstrings",
+    srcs = ["docstrings.py"],
+    visibility = ["//visibility:public"],
+)

tensorflow/python/tools/api/generator2/apis.bzl

@@ -0,0 +1,32 @@
+"""generate_api API definitions."""
+
+load(":patterns.bzl", "compile_patterns")
+
+APIS = {
+    "tf_keras": {
+        "decorator": "tensorflow.python.util.tf_export.keras_export",
+        "target_patterns": compile_patterns([
+            "//third_party/py/tf_keras/...",
+        ]),
+    },
+    "tensorflow": {
+        "decorator": "tensorflow.python.util.tf_export.tf_export",
+        "target_patterns": compile_patterns([
+            "//tensorflow/python/...",
+            "//tensorflow/dtensor/python:accelerator_util",
+            "//tensorflow/dtensor/python:api",
+            "//tensorflow/dtensor/python:config",
+            "//tensorflow/dtensor/python:d_checkpoint",
+            "//tensorflow/dtensor/python:d_variable",
+            "//tensorflow/dtensor/python:input_util",
+            "//tensorflow/dtensor/python:layout",
+            "//tensorflow/dtensor/python:mesh_util",
+            "//tensorflow/dtensor/python:tpu_util",
+            "//tensorflow/dtensor/python:save_restore",
+            "//tensorflow/lite/python/...",
+            "//tensorflow/python:modules_with_exports",
+            "//tensorflow/lite/tools/optimize/debugging/python:all",
+            "//tensorflow/compiler/mlir/quantization/tensorflow/python:all",
+        ]),
+    },
+}

tensorflow/python/tools/api/generator2/docstrings.py

@@ -0,0 +1,63 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Extra docstrings that will be exported for tensorflow APIs."""
+
+# pylint: disable=pointless-string-statement
+"""Import router for absl.app.
+
+API docstring: tensorflow.app
+"""
+
+"""Core module for TensorFlow distribution objects and helpers.
+
+API docstring: tensorflow.distributions
+"""
+
+"""Import router for file_io.
+
+API docstring: tensorflow.gfile
+"""
+
+"""Evaluation-related metrics.
+
+API docstring: tensorflow.metrics
+"""
+
+"""Module for constructing RNN Cells.
+
+API docstring: tensorflow.nn.rnn_cell
+"""
+
+"""Resource management library.
+
+API docstring: tensorflow.resource_loader
+"""
+
+"""System configuration library.
+
+API docstring: tensorflow.sysconfig
+"""
+
+"""Testing.
+
+API docstring: tensorflow.test
+"""
+
+"""Support for training models.
+
+See the [Training](https://tensorflow.org/api_guides/python/train) guide.
+
+API docstring: tensorflow.train
+"""

tensorflow/python/tools/api/generator2/extractor/BUILD

@@ -0,0 +1,42 @@
+load("//tensorflow:pytype.default.bzl", "pytype_strict_binary", "pytype_strict_library")
+load("//tensorflow:tensorflow.bzl", "py_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//tensorflow/python/tools/api/generator2:__subpackages__"],
+    licenses = ["notice"],
+)
+
+pytype_strict_library(
+    name = "extractor",
+    srcs = ["extractor.py"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tensorflow/python/tools/api/generator2/shared:exported_api",
+        "@absl_py//absl/flags",
+        "@absl_py//absl/logging",
+    ],
+)
+
+pytype_strict_binary(
+    name = "main",
+    srcs = ["main.py"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":extractor",
+        "@absl_py//absl:app",
+    ],
+)
+
+py_test(
+    name = "extractor_test",
+    srcs = ["extractor_test.py"],
+    strict_deps = True,
+    tags = ["no_pip"],
+    deps = [
+        ":extractor",
+        "@absl_py//absl/testing:absltest",
+        # copybara:uncomment "//third_party/py/google/protobuf:use_fast_cpp_protos",
+        "//tensorflow/python/tools/api/generator2/shared:exported_api",
+    ],
+)

tensorflow/python/tools/api/generator2/extractor/extractor.py

@@ -0,0 +1,363 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+"""Extracts API information for a set of Python sources."""
+
+import ast
+from collections.abc import Sequence
+import re
+from typing import Any, Optional, Union, cast
+
+from absl import flags
+from absl import logging
+
+from tensorflow.python.tools.api.generator2.shared import exported_api
+
+_OUTPUT = flags.DEFINE_string('output', '', 'File to output contents to.')
+_DECORATOR = flags.DEFINE_string(
+    'decorator',
+    '',
+    'Full path to Python decorator function used for exporting API.',
+)
+_API_NAME = flags.DEFINE_string(
+    'api_name',
+    '',
+    'Prefix for all exported symbols and docstrings.',
+)
+
+_DOCSTRING_PATTERN: re.Pattern[str] = re.compile(
+    r'\s*API\s+docstring:\s*([\w.]+)\s*'
+)
+
+
+class BadExportError(Exception):
+  """Exception for bad exports."""
+
+
+class Parser(ast.NodeVisitor):
+  """Parser for Python source files that extracts TF API exports."""
+
+  _exports: exported_api.ExportedApi
+  _decorator_package: str
+  _decorator_symbol: str
+  _api_name: str
+  _current_file: Optional[str] = None
+  _current_file_decorators: set[str]
+
+  def __init__(
+      self,
+      exports: exported_api.ExportedApi,
+      decorator: str,
+      api_name: str,
+  ):
+    self._exports = exports
+    self._decorator_package, self._decorator_symbol = decorator.rsplit('.', 1)
+    self._api_name = api_name
+
+  def process_file(self, filename: str) -> None:
+    """Finds exported APIs in filename."""
+    try:
+      with open(filename, mode='r', encoding='utf-8') as f:
+        contents = f.read()
+    except Exception as e:  # pylint: disable=broad-exception-caught
+      # log and ignore exceptions from read
+      logging.exception('Error reading %s: %s', filename, e)
+    else:
+      self.process(filename, contents)
+
+  def process(self, filename: str, contents: str) -> None:
+    """Finds exported APIs in contents."""
+    self._current_file_decorators = set()
+    self._current_file = filename
+    try:
+      parsed = ast.parse(contents, filename=filename)
+    except Exception as e:  # pylint: disable=broad-exception-caught
+      # logging errors when parsing file
+      logging.exception('Error parsing %s: %s', filename, e)
+    else:
+      self.visit(parsed)
+    finally:
+      self._current_file = None
+      self._current_file_decorators = set()
+
+  def visit_Module(self, node: ast.Module) -> None:  # pylint: disable=invalid-name
+    for stmt in node.body:
+      self._process_stmt(stmt)
+
+  def _process_stmt(self, node: ast.stmt) -> None:
+    """Process top-level statement for exported apis."""
+    if isinstance(node, (ast.ClassDef, ast.FunctionDef)):
+      self._process_def(node)
+    elif isinstance(node, ast.Assign):
+      self._process_assign(node)
+    elif isinstance(node, ast.Expr):
+      self._process_expr(node)
+    else:
+      self.visit(node)
+
+  def visit_Import(self, node: ast.Import) -> None:  # pylint: disable=invalid-name
+    """Identifies imports of decorator."""
+    for name in node.names:
+      if name.name == self._decorator_package:
+        if name.asname:
+          # import <package> as <name>
+          self._current_file_decorators.add(
+              name.asname + '.' + self._decorator_symbol
+          )
+        else:
+          # import <package>
+          _, module = self._decorator_package.rsplit('.', 1)
+          self._current_file_decorators.add(
+              module + '.' + self._decorator_symbol
+          )
+    self.generic_visit(node)
+
+  def visit_ImportFrom(self, node: ast.ImportFrom) -> None:  # pylint: disable=invalid-name
+    """Identifies imports of decorator."""
+    if node.module == self._decorator_package:
+      for name in node.names:
+        if name.name == self._decorator_symbol:
+          if name.asname:
+            # from <package> import <symbol> as <name>
+            self._current_file_decorators.add(name.asname)
+          else:
+            # from <package> import <symbol>
+            self._current_file_decorators.add(name.name)
+    else:
+      parent, module = self._decorator_package.rsplit('.', 1)
+      if node.module == parent:
+        for name in node.names:
+          if name.name == module:
+            if name.asname:
+              # from <parent> import <module> as <name>
+              self._current_file_decorators.add(
+                  name.asname + '.' + self._decorator_symbol
+              )
+            else:
+              # from <parent> import <module>
+              self._current_file_decorators.add(
+                  name.name + '.' + self._decorator_symbol
+              )
+    self.generic_visit(node)
+
+  def _process_def(self, node: Union[ast.ClassDef, ast.FunctionDef]) -> None:
+    """Process top-level [Class|Function]Def for potential symbol export."""
+    # @tf_export(...)
+    # [class|def] <id>:
+    for decorator in node.decorator_list:
+      if self._is_export_call(decorator):
+        self._add_exported_symbol(cast(ast.Call, decorator), node.name)
+      else:
+        self.visit(decorator)
+
+    if isinstance(node, ast.ClassDef):
+      for base in node.bases:
+        self.visit(base)
+      for kw in node.keywords:
+        self.visit(kw)
+    elif isinstance(node, ast.FunctionDef):
+      self.visit(node.args)
+      if node.returns:
+        self.visit(node.returns)
+
+    for stmt in node.body:
+      self.visit(stmt)
+
+  def _process_assign(self, node: ast.Assign) -> None:
+    """Process top-level assign for potential symbol export."""
+    if isinstance(node.value, ast.Call) and self._is_export_call(
+        node.value.func
+    ):
+      # id = tf_export(...)(...)
+      if len(node.targets) != 1:
+        raise BadExportError(
+            f'{self._current_file}:{node.lineno} export must be'
+            f' assigned to a single value: {ast.dump(node)}'
+        )
+      symbol = self._name(node.targets[0])
+      if not symbol:
+        raise BadExportError(
+            f'{self._current_file}:{node.lineno} export must be'
+            f' assigned to a single value: {ast.dump(node)}'
+        )
+      self._add_exported_symbol(node.value.func, symbol)
+    else:
+      self.visit(node)
+
+  def _process_expr(self, node: ast.Expr) -> None:
+    """Process top-level expression for potential symbol export."""
+    if isinstance(node.value, ast.Call):
+      self._process_call(node.value)
+    elif isinstance(node.value, ast.Constant):
+      self._process_constant(node.value)
+    else:
+      self.visit(node)
+
+  def _process_call(self, node: ast.Call) -> None:
+    """Process top-level call for potential symbol export."""
+    func = node.func
+    if self._is_export_call(func):
+      func = cast(ast.Call, func)
+      # tf_export(...)(id)
+      if len(node.args) != 1 or node.keywords:
+        raise BadExportError(
+            f'{self._current_file}:{node.lineno} export must be'
+            f' called with a single value: {ast.dump(node)}'
+        )
+      symbol = self._name(self._unwrap_simple_call(node.args[0]))
+      if not symbol:
+        raise BadExportError(
+            f'{self._current_file}:{node.lineno} export must be'
+            f' called with a single value: {ast.dump(node)}'
+        )
+      self._add_exported_symbol(func, symbol)
+    elif (
+        isinstance(func, ast.Attribute)
+        and func.attr == 'export_constant'
+        and self._is_export_call(func.value)
+    ):
+      # tf_export(...).export_constant(__name__, id)
+      if (
+          len(node.args) != 2
+          or node.keywords
+          or self._name(node.args[0]) != '__name__'
+      ):
+        raise BadExportError(
+            f'{self._current_file}:{node.lineno} export_constant must be'
+            f' called with __name__, <id>: {ast.dump(node)}'
+        )
+      self._add_exported_symbol(func.value, self._literal_value(node.args[1]))
+    else:
+      self.visit(node)
+
+  def _process_constant(self, node: ast.Constant) -> None:
+    """Process top-level constant for a potential API docstring export."""
+    if isinstance(node.value, str):
+      docstring, modules = self._extract_docstring(node.value)
+      if modules:
+        self._exports.add_doc(
+            exported_api.ExportedDoc.create(
+                file_name=self._current_file,
+                line_no=node.lineno,
+                modules=modules,
+                docstring=docstring,
+            )
+        )
+      else:
+        self.visit(node)
+
+  def _extract_docstring(self, value: str) -> tuple[str, Sequence[str]]:
+    """Extract docstring and list of modules that it should be applied to."""
+    docstring = ''
+    modules = []
+    for line in value.splitlines():
+      match = _DOCSTRING_PATTERN.match(line)
+      if match:
+        module = match.group(1).strip()
+        # API docstring: <module>
+        if module == self._api_name or module.startswith(self._api_name + '.'):
+          modules.append(module)
+      else:
+        docstring += line + '\n'
+    return (docstring.strip(), modules)
+
+  def visit_Call(self, node: ast.Call) -> None:  # pylint: disable=invalid-name
+    if self._is_export_call(node):
+      raise BadExportError(
+          f'{self._current_file}:{node.lineno} export must be'
+          f' used at top level of file: {ast.dump(node)}'
+      )
+    self.generic_visit(node)
+
+  def visit_Constant(self, node: ast.Constant) -> None:
+    if isinstance(node.value, str):
+      _, modules = self._extract_docstring(node.value)
+      if modules:
+        raise BadExportError(
+            f'{self._current_file}:{node.lineno} API docstrings must be'
+            f' at top level of file: {ast.dump(node)}'
+        )
+    self.generic_visit(node)
+
+  def _is_export_call(self, node: ast.expr) -> bool:  # TypeGuard[ast.Call]
+    return (
+        isinstance(node, ast.Call)
+        and self._name(node.func) in self._current_file_decorators
+    )
+
+  def _name(self, node: ast.expr) -> Optional[str]:
+    if isinstance(node, ast.Name):
+      return node.id
+    if isinstance(node, ast.Attribute):
+      parent = self._name(node.value)
+      if parent:
+        return f'{parent}.{node.attr}'
+
+  def _unwrap_simple_call(self, node: ast.expr) -> ast.expr:
+    """Unwraps a function call that takes a single unnamed parameter."""
+    if isinstance(node, ast.Call) and len(node.args) == 1 and not node.keywords:
+      return self._unwrap_simple_call(node.args[0])
+    return node
+
+  def _literal_value(self, node: ast.expr) -> Any:
+    try:
+      return ast.literal_eval(node)
+    except Exception as e:
+      raise BadExportError(
+          f'{self._current_file}:{node.lineno} all arguments to'
+          f' export must be literal values: {ast.dump(node)}'
+      ) from e
+
+  def _add_exported_symbol(self, node: ast.Call, symbol_name: str) -> None:
+    """Adds an exported symbol represented by the given call."""
+    if symbol_name.find('.') != -1:
+      raise BadExportError(
+          f'{self._current_file}:{node.lineno} export called with symbol'
+          f' {symbol_name} not defined in current file: {ast.dump(node)}'
+      )
+    v2_apis = tuple(
+        f'{self._api_name}.{self._literal_value(arg)}' for arg in node.args
+    )
+    v1_apis = v2_apis
+    for kw in node.keywords:
+      if kw.arg == 'v1':
+        v1_apis = tuple(
+            f'{self._api_name}.{v}' for v in self._literal_value(kw.value)
+        )
+      elif kw.arg == 'allow_multiple_exports':
+        # no-op kept for backward compatibility of `tf-keras` with TF 2.13
+        pass
+      else:
+        raise BadExportError(
+            f'{self._current_file}:{node.lineno} export called'
+            f' with unknown argument {kw.arg}: {ast.dump(node)}'
+        )
+    self._exports.add_symbol(
+        exported_api.ExportedSymbol.create(
+            file_name=self._current_file,
+            line_no=node.lineno,
+            symbol_name=symbol_name,
+            v2_apis=v2_apis,
+            v1_apis=v1_apis,
+        )
+    )
+
+
+def main(argv: Sequence[str]) -> None:
+  exporter = exported_api.ExportedApi()
+  p = Parser(exporter, _DECORATOR.value, _API_NAME.value)
+  for arg in argv[1:]:
+    p.process_file(arg)
+
+  exporter.write(_OUTPUT.value)

tensorflow/python/tools/api/generator2/extractor/extractor_test.py

@@ -0,0 +1,256 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+from absl.testing import absltest
+
+from tensorflow.python.tools.api.generator2.extractor import extractor
+from tensorflow.python.tools.api.generator2.shared import exported_api
+
+
+class ParserTest(absltest.TestCase):
+
+  def test_exported_docstring(self):
+    exporter = exported_api.ExportedApi()
+    p = extractor.Parser(
+        exporter,
+        decorator='tf.tf_export',
+        api_name='tf',
+    )
+    p.process(
+        'test.py',
+        '''# 1
+"""this is an exported docstring.
+API docstring: tf.test
+"""  # 4
+        ''',
+    )
+    self.assertEqual(
+        exporter,
+        exported_api.ExportedApi(
+            docs=[
+                exported_api.ExportedDoc(
+                    file_name='test.py',
+                    line_no=2,
+                    docstring='this is an exported docstring.',
+                    modules=('tf.test',),
+                )
+            ],
+        ),
+    )
+
+  def test_exported_docstring_not_at_top_level(self):
+    exporter = exported_api.ExportedApi()
+    p = extractor.Parser(
+        exporter,
+        decorator='tf.tf_export',
+        api_name='tf',
+    )
+    self.assertRaisesRegex(
+        extractor.BadExportError,
+        'test.py:3',
+        lambda: p.process(  # pylint: disable=g-long-lambda
+            'test.py',
+            '''# 1
+def a():  # 2
+  """a docstring
+  API docstring: tf.test
+  """  # 5
+          ''',
+        ),
+    )
+
+  def test_exported_symbol(self):
+    exporter = exported_api.ExportedApi()
+    p = extractor.Parser(
+        exporter,
+        decorator='extractor.api_export.tf_export',
+        api_name='tf',
+    )
+    p.process(
+        'test.py',
+        """# 1
+from extractor import api_export  # 2
+from extractor import api_export as ae  # 3
+try:  # 4
+  from extractor.api_export import tf_export  # 5
+except ImportError:  # 6
+  pass  # 7
+from extractor.api_export import tf_export as tfe # 8
+from extractor.api_export import other_export  # 9
+_a = api_export.tf_export("a")(foo)  # 10
+api_export.tf_export("b", v1=["v1_b"])(_b)  # 11
+tfe("c")(_c)  # 12
+@ae.tf_export("d")  # 13
+class _D():  # 14
+  pass  # 15
+@api_export.tf_export("e", "e_v2", v1=[])  # 16
+def _e():  # 17
+  pass  # 18
+tf_export(v1=["f", "f_alias"])(  # 19
+    dispatch.dispatch(deprecation(_f))  # 20
+)  # 21
+@other_export("not-exported")  # 22
+def _not_exported():  # 23
+  pass  # 24
+        """,
+    )
+    self.assertEqual(
+        exporter,
+        exported_api.ExportedApi(
+            symbols=[
+                exported_api.ExportedSymbol(
+                    file_name='test.py',
+                    line_no=10,
+                    symbol_name='_a',
+                    v1_apis=('tf.a',),
+                    v2_apis=('tf.a',),
+                ),
+                exported_api.ExportedSymbol(
+                    file_name='test.py',
+                    line_no=11,
+                    symbol_name='_b',
+                    v1_apis=('tf.v1_b',),
+                    v2_apis=('tf.b',),
+                ),
+                exported_api.ExportedSymbol(
+                    file_name='test.py',
+                    line_no=12,
+                    symbol_name='_c',
+                    v1_apis=('tf.c',),
+                    v2_apis=('tf.c',),
+                ),
+                exported_api.ExportedSymbol(
+                    file_name='test.py',
+                    line_no=13,
+                    symbol_name='_D',
+                    v1_apis=('tf.d',),
+                    v2_apis=('tf.d',),
+                ),
+                exported_api.ExportedSymbol(
+                    file_name='test.py',
+                    line_no=16,
+                    symbol_name='_e',
+                    v1_apis=(),
+                    v2_apis=('tf.e', 'tf.e_v2'),
+                ),
+                exported_api.ExportedSymbol(
+                    file_name='test.py',
+                    line_no=19,
+                    symbol_name='_f',
+                    v1_apis=('tf.f', 'tf.f_alias'),
+                    v2_apis=(),
+                ),
+            ],
+        ),
+    )
+
+  def test_exported_symbol_not_at_top_level(self):
+    exporter = exported_api.ExportedApi()
+    p = extractor.Parser(
+        exporter,
+        decorator='tf.tf_export',
+        api_name='tf',
+    )
+    self.assertRaisesRegex(
+        extractor.BadExportError,
+        'test.py:4',
+        lambda: p.process(  # pylint: disable=g-long-lambda
+            'test.py',
+            """# 1
+from tf import tf_export  # 2
+def method():  # 3
+  tf_export("a")(a)  # 4
+            """,
+        ),
+    )
+
+  def test_exported_symbol_not_applied(self):
+    exporter = exported_api.ExportedApi()
+    p = extractor.Parser(
+        exporter,
+        decorator='tf.tf_export',
+        api_name='tf',
+    )
+    self.assertRaisesRegex(
+        extractor.BadExportError,
+        'test.py:3',
+        lambda: p.process(  # pylint: disable=g-long-lambda
+            'test.py',
+            """# 1
+from tf import tf_export  # 2
+tf_export("a")  # 3
+            """,
+        ),
+    )
+
+  def test_exported_symbol_non_literal_args(self):
+    exporter = exported_api.ExportedApi()
+    p = extractor.Parser(
+        exporter,
+        decorator='tf.tf_export',
+        api_name='tf',
+    )
+    self.assertRaisesRegex(
+        extractor.BadExportError,
+        'test.py:3',
+        lambda: p.process(  # pylint: disable=g-long-lambda
+            'test.py',
+            """# 1
+from tf import tf_export  # 2
+tf_export(a)(b)  # 3
+            """,
+        ),
+    )
+
+  def test_exported_symbol_unknown_args(self):
+    exporter = exported_api.ExportedApi()
+    p = extractor.Parser(
+        exporter,
+        decorator='tf.tf_export',
+        api_name='tf',
+    )
+    self.assertRaisesRegex(
+        extractor.BadExportError,
+        'test.py:3',
+        lambda: p.process(  # pylint: disable=g-long-lambda
+            'test.py',
+            """# 1
+from tf import tf_export  # 2
+tf_export(a)(b)  # 3
+            """,
+        ),
+    )
+
+  def test_exported_symbol_includes_module(self):
+    exporter = exported_api.ExportedApi()
+    p = extractor.Parser(
+        exporter,
+        decorator='tf.tf_export',
+        api_name='tf',
+    )
+    self.assertRaisesRegex(
+        extractor.BadExportError,
+        'test.py:3',
+        lambda: p.process(  # pylint: disable=g-long-lambda
+            'test.py',
+            """# 1
+from tf import tf_export  # 2
+tf_export(a)(x.b)  # 3
+            """,
+        ),
+    )
+
+
+if __name__ == '__main__':
+  absltest.main()

tensorflow/python/tools/api/generator2/extractor/main.py

@@ -0,0 +1,22 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+"""Entrypoint for API Extractor."""
+
+from absl import app
+
+from tensorflow.python.tools.api.generator2.extractor import extractor
+
+if __name__ == '__main__':
+  app.run(extractor.main)

tensorflow/python/tools/api/generator2/generate_api.bzl

@@ -0,0 +1,398 @@
+"""Rules to generate the TensorFlow public API from annotated files."""
+
+load("@bazel_skylib//lib:paths.bzl", "paths")
+load("@rules_python//python:py_info.bzl", "PyInfo")
+load("//tensorflow/python/tools/api/generator:api_init_files.bzl", "TENSORFLOW_API_INIT_FILES")
+load(":apis.bzl", _APIS = "APIS")
+load(":patterns.bzl", "any_match")
+
+APIS = _APIS.keys()
+
+_MODULE_PREFIX = ""
+
+def _api_info_init(*, transitive_api):
+    if type(transitive_api) != type(depset()):
+        fail("ApiInfo.transitive_api must be a depset")
+    return {"transitive_api": transitive_api}
+
+ApiInfo, _new_api_info = provider(
+    doc = "Provider for API symbols and docstrings extracted from Python files.",
+    fields = {
+        "transitive_api": "depset of files with extracted API.",
+    },
+    init = _api_info_init,
+)
+
+def _py_files(f):
+    if f.basename.endswith(".py") or f.basename.endswith(".py3"):
+        return f.path
+    return None
+
+def _merge_py_info(
+        deps,
+        direct_sources = None,
+        direct_imports = None,
+        has_py2_only_sources = False,
+        has_py3_only_sources = False,
+        uses_shared_libraries = False):
+    transitive_sources = []
+    transitive_imports = []
+    for dep in deps:
+        if PyInfo in dep:
+            transitive_sources.append(dep[PyInfo].transitive_sources)
+            transitive_imports.append(dep[PyInfo].imports)
+            has_py2_only_sources = has_py2_only_sources or dep[PyInfo].has_py2_only_sources
+            has_py3_only_sources = has_py3_only_sources or dep[PyInfo].has_py3_only_sources
+            uses_shared_libraries = uses_shared_libraries or dep[PyInfo].uses_shared_libraries
+
+    return PyInfo(
+        transitive_sources = depset(direct = direct_sources, transitive = transitive_sources),
+        imports = depset(direct = direct_imports, transitive = transitive_imports),
+        has_py2_only_sources = has_py2_only_sources,
+        has_py3_only_sources = has_py3_only_sources,
+        uses_shared_libraries = uses_shared_libraries,
+    )
+
+def _merge_api_info(
+        deps,
+        direct_api = None):
+    transitive_api = []
+    for dep in deps:
+        if ApiInfo in dep:
+            transitive_api.append(dep[ApiInfo].transitive_api)
+    return ApiInfo(transitive_api = depset(direct = direct_api, transitive = transitive_api))
+
+def _api_extractor_impl(target, ctx):
+    api = ctx.attr.api
+    config = _APIS[api]
+    direct_api = []
+
+    # Make sure the rule has a non-empty srcs attribute.
+    if (
+        any_match(config["target_patterns"], target.label) and
+        hasattr(ctx.rule.attr, "srcs") and
+        ctx.rule.attr.srcs
+    ):
+        output = ctx.actions.declare_file("_".join([
+            target.label.name,
+            "extracted",
+            api,
+            "api.json",
+        ]))
+
+        args = ctx.actions.args()
+        args.set_param_file_format("multiline")
+        args.use_param_file("--flagfile=%s")
+
+        args.add("--output", output)
+        args.add("--decorator", config["decorator"])
+        args.add("--api_name", api)
+        args.add_all(ctx.rule.files.srcs, expand_directories = True, map_each = _py_files)
+
+        ctx.actions.run(
+            mnemonic = "ExtractAPI",
+            executable = ctx.executable._extractor_bin,
+            inputs = ctx.rule.files.srcs,
+            outputs = [output],
+            arguments = [args],
+            progress_message = "Extracting " + api + " APIs for %{label} to %{output}.",
+            use_default_shell_env = True,
+        )
+
+        direct_api.append(output)
+
+    return [
+        _merge_api_info(ctx.rule.attr.deps if hasattr(ctx.rule.attr, "deps") else [], direct_api = direct_api),
+    ]
+
+api_extractor = aspect(
+    doc = "Extracts the exported API for the given target and its dependencies.",
+    implementation = _api_extractor_impl,
+    attr_aspects = ["deps"],
+    provides = [ApiInfo],
+    # Currently the Python rules do not correctly advertise their providers.
+    # required_providers = [PyInfo],
+    attrs = {
+        "_extractor_bin": attr.label(
+            default = Label("//tensorflow/python/tools/api/generator2/extractor:main"),
+            executable = True,
+            cfg = "exec",
+        ),
+        "api": attr.string(
+            doc = "API to extract from dependencies.",
+            mandatory = True,
+            values = APIS,
+        ),
+    },
+)
+
+def _extract_api_impl(ctx):
+    return [
+        _merge_api_info(ctx.attr.deps),
+        _merge_py_info(ctx.attr.deps),
+    ]
+
+extract_api = rule(
+    doc = "Extract Python API for all targets in transitive dependencies.",
+    implementation = _extract_api_impl,
+    attrs = {
+        "deps": attr.label_list(
+            doc = "Targets to extract API from.",
+            allow_empty = False,
+            aspects = [api_extractor],
+            providers = [PyInfo],
+            mandatory = True,
+        ),
+        "api": attr.string(
+            doc = "API to extract from dependencies.",
+            mandatory = True,
+            values = APIS,
+        ),
+    },
+    provides = [ApiInfo, PyInfo],
+)
+
+def _get_module_by_path(dir_path, output_dir):
+    """Get module that corresponds to the path.
+
+    bazel-out/k8-opt/bin/tensorflow/_api/v2/compat/v2/compat/v2/compat/__init__.py
+    to
+    tensorflow._api.v2.compat.v2.compat.v2.compat
+
+    Args:
+    dir_path: Path to the directory.
+    output_dir: Path to the directory.
+
+    Returns:
+    Name of module that corresponds to the given directory.
+    """
+    dir_path = dir_path.split(output_dir)[1]
+    dir_path = dir_path.replace("__init__.py", "")
+
+    return dir_path.replace("/", ".").strip(".")
+
+def _generate_api_impl(ctx):
+    args = ctx.actions.args()
+    args.set_param_file_format("multiline")
+    args.use_param_file("--flagfile=%s")
+
+    args.add_joined("--output_files", ctx.outputs.output_files, join_with = ",")
+    args.add("--output_dir", paths.join(ctx.bin_dir.path, ctx.label.workspace_root, ctx.label.package, ctx.attr.output_dir))
+    if ctx.file.root_init_template:
+        args.add("--root_init_template", ctx.file.root_init_template)
+    args.add("--apiversion", ctx.attr.api_version)
+    args.add_joined("--compat_api_versions", ctx.attr.compat_api_versions, join_with = ",")
+    args.add_joined("--compat_init_templates", ctx.files.compat_init_templates, join_with = ",")
+    args.add("--output_package", ctx.attr.output_package)
+    args.add_joined("--packages_to_ignore", ctx.attr.packages_to_ignore, join_with = ",")
+    if _MODULE_PREFIX:
+        args.add("--module_prefix", _MODULE_PREFIX)
+    if ctx.attr.use_lazy_loading:
+        args.add("--use_lazy_loading")
+    else:
+        args.add("--nouse_lazy_loading")
+    if ctx.attr.proxy_module_root:
+        args.add("--proxy_module_root", ctx.attr.proxy_module_root)
+    args.add_joined("--file_prefixes_to_strip", [ctx.bin_dir.path, ctx.genfiles_dir.path] + ctx.attr.file_prefixes_to_strip, join_with = ",")
+    if ctx.attr.root_file_name:
+        args.add("--root_file_name", ctx.attr.root_file_name)
+
+    inputs = depset(transitive = [
+        dep[ApiInfo].transitive_api
+        for dep in ctx.attr.deps
+    ])
+    args.add_all(
+        inputs,
+        expand_directories = True,
+    )
+
+    transitive_inputs = [inputs]
+    if ctx.attr.root_init_template:
+        transitive_inputs.append(ctx.attr.root_init_template.files)
+
+    ctx.actions.run(
+        mnemonic = "GenerateAPI",
+        executable = ctx.executable._generator_bin,
+        inputs = depset(
+            direct = ctx.files.compat_init_templates,
+            transitive = transitive_inputs,
+        ),
+        outputs = ctx.outputs.output_files,
+        arguments = [args],
+        progress_message = "Generating APIs for %{label} to %{output}.",
+        use_default_shell_env = True,
+    )
+
+    # Convert output_paths to the list of corresponding modules for the further testing
+    if ctx.outputs.api_packages_path:
+        output_modules = sorted([
+            _get_module_by_path(f.path, ctx.bin_dir.path)
+            for f in ctx.outputs.output_files
+            if "__init__.py" in f.path
+        ])
+        ctx.actions.write(ctx.outputs.api_packages_path, "\n".join(output_modules))
+
+generate_api = rule(
+    doc = "Generate Python API for all targets in transitive dependencies.",
+    implementation = _generate_api_impl,
+    attrs = {
+        "deps": attr.label_list(
+            doc = "extract_api targets to generate API from.",
+            allow_empty = True,
+            providers = [ApiInfo, PyInfo],
+            mandatory = True,
+        ),
+        "root_init_template": attr.label(
+            doc = "Template for the top level __init__.py file",
+            allow_single_file = True,
+        ),
+        "api_packages_path": attr.output(
+            doc = "Name of the file with the list of all API packages.",
+        ),
+        "api_version": attr.int(
+            doc = "The API version to generate (1 or 2)",
+            values = [1, 2],
+        ),
+        "compat_api_versions": attr.int_list(
+            doc = "Additional versions to generate in compat/ subdirectory.",
+        ),
+        "compat_init_templates": attr.label_list(
+            doc = "Template for top-level __init__files under compat modules. This list must be " +
+                  "in the same order as the list of versions in compat_apiversions",
+            allow_files = True,
+        ),
+        "output_package": attr.string(
+            doc = "Root output package.",
+        ),
+        "output_dir": attr.string(
+            doc = "Subdirectory to output API to. If non-empty, must end with '/'.",
+        ),
+        "proxy_module_root": attr.string(
+            doc = "Module root for proxy-import format. If specified, proxy files with " +
+                  "`from proxy_module_root.proxy_module import *` will be created to enable " +
+                  "import resolution under TensorFlow.",
+        ),
+        "output_files": attr.output_list(
+            doc = "List of __init__.py files that should be generated. This list should include " +
+                  "file name for every module exported using tf_export. For e.g. if an op is " +
+                  "decorated with @tf_export('module1.module2', 'module3'). Then, output_files " +
+                  "should include module1/module2/__init__.py and module3/__init__.py.",
+        ),
+        "use_lazy_loading": attr.bool(
+            doc = "If true, lazy load imports in the generated API rather then importing them all statically.",
+        ),
+        "packages_to_ignore": attr.string_list(
+            doc = "List of packages to ignore tf_exports from.",
+        ),
+        "root_file_name": attr.string(
+            doc = "The file name that should be generated for the top level API.",
+        ),
+        "file_prefixes_to_strip": attr.string_list(
+            doc = "The file prefixes to strip from the import paths. Ex: bazel's bin and genfile",
+        ),
+        "_generator_bin": attr.label(
+            default = Label("//tensorflow/python/tools/api/generator2/generator:main"),
+            executable = True,
+            cfg = "exec",
+        ),
+    },
+)
+
+def generate_apis(
+        name,
+        apis = ["tensorflow"],
+        deps = [
+            "//tensorflow/python:no_contrib",
+            "//tensorflow/python:modules_with_exports",
+            "//tensorflow/lite/python:analyzer",
+            "//tensorflow/lite/python:lite",
+            "//tensorflow/lite/python/authoring",
+        ],
+        output_files = TENSORFLOW_API_INIT_FILES,
+        root_init_template = None,
+        api_packages_file_name = None,
+        api_version = 2,
+        compat_api_versions = [],
+        compat_init_templates = [],
+        output_package = "tensorflow",
+        output_dir = "",
+        proxy_module_root = None,
+        packages_to_ignore = [],
+        root_file_name = None,
+        visibility = ["//visibility:private"],
+        file_prefixes_to_strip = []):
+    """Generate TensorFlow APIs for a set of libraries.
+
+    Args:
+        name: name of generate_api target.
+        apis: APIs to extract. See APIS constant for allowed values.
+        deps: python_library targets to serve as roots for extracting APIs.
+        output_files: The list of files that the API generator is expected to create.
+        root_init_template: The template for the top level __init__.py file generated.
+            "#API IMPORTS PLACEHOLDER" comment will be replaced with imports.
+        api_packages_file_name: Name of the file with the list of all API packages. Stores in output_dir.
+        api_version: THhe API version to generate. (1 or 2)
+        compat_api_versions: Additional versions to generate in compat/ subdirectory.
+        compat_init_templates: Template for top level __init__.py files under the compat modules.
+            The list must be in the same order as the list of versions in 'compat_api_versions'
+        output_package: Root output package.
+        output_dir: Directory where the generated output files are placed. This should be a prefix
+            of every directory in 'output_files'
+        proxy_module_root: Module root for proxy-import format. If specified, proxy files with
+            `from proxy_module_root.proxy_module import *` will be created to enable import
+            resolution under TensorFlow.
+        packages_to_ignore: List of packages to ignore tf_exports from.
+        root_file_name: The file name that should be generated for the top level API.
+        visibility: Visibility of the target containing the generated files.
+    """
+    extract_api_targets = []
+    for api in apis:
+        extract_name = name + ".extract-" + api
+        extract_api(
+            name = extract_name,
+            api = api,
+            deps = deps,
+            visibility = ["//visibility:private"],
+        )
+        extract_api_targets.append(extract_name)
+
+    if root_file_name != None and root_file_name != "__init__.py":
+        # Rename file for top-level API.
+        output_files = [root_file_name if f == "__init__.py" else f for f in output_files]
+    elif proxy_module_root != None:
+        # Avoid conflicts between the __init__.py file of TensorFlow and proxy module.
+        output_files = [f for f in output_files if f != "__init__.py"]
+
+    all_output_files = [paths.join(output_dir, f) for f in output_files]
+
+    if api_packages_file_name:
+        api_packages_path = "%s%s" % (output_dir, api_packages_file_name)
+    else:
+        api_packages_path = None
+
+    generate_api(
+        name = name,
+        deps = extract_api_targets,
+        output_files = all_output_files,
+        output_dir = output_dir,
+        root_init_template = root_init_template,
+        compat_api_versions = compat_api_versions,
+        compat_init_templates = compat_init_templates,
+        api_version = api_version,
+        proxy_module_root = proxy_module_root,
+        visibility = visibility,
+        packages_to_ignore = packages_to_ignore,
+        # copybara:uncomment_begin(configurable API loading)
+        # use_lazy_loading = select({
+        # "//tensorflow/python/tools/api/generator2:static_gen": False,
+        # "//tensorflow:api_indexable": False,
+        # "//conditions:default": True,
+        # }),
+        # copybara:uncomment_end_and_comment_begin
+        use_lazy_loading = False,
+        # copybara:comment_end
+        output_package = output_package,
+        root_file_name = root_file_name,
+        api_packages_path = api_packages_path,
+        file_prefixes_to_strip = file_prefixes_to_strip,
+    )

tensorflow/python/tools/api/generator2/generator/BUILD

@@ -0,0 +1,43 @@
+load("//tensorflow:pytype.default.bzl", "pytype_strict_binary", "pytype_strict_library")
+load("//tensorflow:tensorflow.bzl", "py_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//tensorflow/python/tools/api/generator2:__subpackages__"],
+    licenses = ["notice"],
+)
+
+pytype_strict_library(
+    name = "generator",
+    srcs = ["generator.py"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tensorflow/python/tools/api/generator2/shared:exported_api",
+        "@absl_py//absl:app",
+        "@absl_py//absl/flags",
+    ],
+)
+
+pytype_strict_binary(
+    name = "main",
+    srcs = ["main.py"],
+    visibility = ["//visibility:public"],
+    deps = [
+        ":generator",
+        "@absl_py//absl:app",
+    ],
+)
+
+py_test(
+    name = "generator_test",
+    srcs = ["generator_test.py"],
+    strict_deps = True,
+    tags = ["no_pip"],
+    deps = [
+        ":generator",
+        "@absl_py//absl/testing:absltest",
+        "@absl_py//absl/testing:parameterized",
+        # copybara:uncomment "//third_party/py/google/protobuf:use_fast_cpp_protos",
+        "//tensorflow/python/tools/api/generator2/shared:exported_api",
+    ],
+)

tensorflow/python/tools/api/generator2/generator/generator.py

@@ -0,0 +1,761 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+"""Library that generates the API for tensorflow."""
+
+import collections
+from collections.abc import Mapping, Sequence, Set
+import dataclasses
+import os
+from typing import Optional
+
+from absl import app
+from absl import flags
+
+from tensorflow.python.tools.api.generator2.shared import exported_api
+
+_OUTPUT_FILES = flags.DEFINE_list(
+    'output_files', None, 'List of files expected to generate.'
+)
+_OUTPUT_DIR = flags.DEFINE_string(
+    'output_dir',
+    None,
+    'Directory where the generated output files are placed. This should be a'
+    ' prefix of every directory in "output_files".',
+)
+_ROOT_INIT_TEMPLATE = flags.DEFINE_string(
+    'root_init_template',
+    None,
+    'Template for top level __init__.py file.  "#API IMPORTS PLACEHOLDER"'
+    ' comment will be replaced with imports.',
+)
+_API_VERSION = flags.DEFINE_integer(
+    'apiversion', 2, 'The API version to generate. (1 or 2)'
+)
+_COMPAT_API_VERSIONS = flags.DEFINE_list(
+    'compat_api_versions',
+    [],
+    'Additional versions to generate in compat/ subdirectory.',
+)
+_COMPAT_INIT_TEMPLATES = flags.DEFINE_list(
+    'compat_init_templates',
+    [],
+    'Template for top-level __init__.py files under compat modules. This list'
+    ' must be in the same order as the list of versions in'
+    ' "compat_apiversions".',
+)
+_OUTPUT_PACKAGE = flags.DEFINE_string(
+    'output_package', 'tensorflow', 'Root output package.'
+)
+_USE_LAZY_LOADING = flags.DEFINE_bool(
+    'use_lazy_loading',
+    True,
+    'If true, lazily load imports rather than loading them all in the'
+    ' __init__.py files. Defaults to true.',
+)
+_PROXY_MODULE_ROOT = flags.DEFINE_string(
+    'proxy_module_root',
+    None,
+    'Module root for proxy-import format. If specified, proxy files with `from'
+    ' proxy_module_root.proxy_module import *` will be created to enable import'
+    ' resolution under TensorFlow.',
+)
+_FILE_PREFIXES_TO_STRIP = flags.DEFINE_list(
+    'file_prefixes_to_strip',
+    [],
+    "File prefixes to strip from the import paths. Ex: bazel's bin and genfile"
+    ' directories.',
+)
+_PACKAGES_TO_IGNORE = flags.DEFINE_list(
+    'packages_to_ignore',
+    [],
+    'Comma separated list of packages to ignore tf_exports from. Ex:'
+    ' packages_to_ignore="tensorflow.python.framework.test_ops"'
+    ' will not export any tf_exports from test_ops',
+)
+_MODULE_PREFIX = flags.DEFINE_string(
+    'module_prefix', '', 'Prefix to append to all imported modules.'
+)
+_ROOT_FILE_PATH = flags.DEFINE_string(
+    'root_file_name',
+    '__init__.py',
+    'The file name that should be generated for the top level API.',
+)
+
+_GENERATED_FILE_HEADER = """# This file is MACHINE GENERATED! Do not edit.
+# Generated by: tensorflow/python/tools/api/generator2/generator/generator.py script.
+\"\"\"%s
+\"\"\"
+
+import sys as _sys
+
+"""
+
+_LAZY_LOADING_MODULE_TEXT_TEMPLATE = """
+# Inform pytype that this module is dynamically populated (b/111239204).
+_HAS_DYNAMIC_ATTRIBUTES = True
+_PUBLIC_APIS = {
+%s
+}
+"""
+_DEPRECATION_FOOTER = """
+from tensorflow.python.util import module_wrapper as _module_wrapper
+
+if not isinstance(_sys.modules[__name__], _module_wrapper.TFModuleWrapper):
+  _sys.modules[__name__] = _module_wrapper.TFModuleWrapper(
+      _sys.modules[__name__], "%s", public_apis=%s, deprecation=%s,
+      has_lite=%s)
+"""
+
+
+class DocExportedTwiceError(Exception):
+  """Exception for when two docstrings are registered to a single module."""
+
+
+def _get_import_path(
+    file: str, file_prefixes_to_strip: Sequence[str], module_prefix: str
+) -> str:
+  module_import_path = file
+  for prefix in file_prefixes_to_strip:
+    module_import_path = module_import_path.removeprefix(prefix)
+  module_import_path = module_import_path.removesuffix('.py')
+  module_import_path = module_import_path.removesuffix('__init__')
+  module_import_path = module_import_path.strip('/')
+  module_import_path = module_import_path.replace('/', '.')
+
+  return module_prefix + module_import_path
+
+
+@dataclasses.dataclass(frozen=True)
+class _Entrypoint:
+  """An entrypoint that was exposed by the use of a decorator.
+
+  Attributes:
+    module: The public module that the symbol was exposed to. For example:
+      tensorflow.io.
+    name: The name the symbol was exported as. For example: decode_png.
+    exported_symbol: The symbol that this entrypoint refers back to.
+  """
+
+  module: str
+  name: str
+  exported_symbol: exported_api.ExportedSymbol
+
+  def get_import(
+      self,
+      file_prefixes_to_strip: Sequence[str],
+      module_prefix: str,
+      use_lazy_loading: bool,
+  ) -> str:
+    """Returns the import statement for this entrypoint.
+
+    Args:
+      file_prefixes_to_strip: List of prefixes to strip from the file name.
+      module_prefix: A prefix to add to the import.
+      use_lazy_loading: Whether to use lazy loading or not.
+    """
+    module_import_path = _get_import_path(
+        self.exported_symbol.file_name, file_prefixes_to_strip, module_prefix
+    )
+    alias = ''
+    symbol_name = self.exported_symbol.symbol_name
+    if self.name != symbol_name:
+      alias = f' as {self.name}'
+    if not use_lazy_loading:
+      return (
+          f'from {module_import_path} import'
+          f' {symbol_name}{alias} # line:'
+          f' {self.exported_symbol.line_no}'
+      )
+    else:
+      return (
+          f"  '{self.name}': ('{module_import_path}',"
+          f" '{symbol_name}'), # line:"
+          f' {self.exported_symbol.line_no}'
+      )
+
+
+@dataclasses.dataclass(frozen=True)
+class PublicAPI:
+  v1_entrypoints_by_module: Mapping[str, set[_Entrypoint]]
+  v2_entrypoints_by_module: Mapping[str, set[_Entrypoint]]
+  v1_generated_imports_by_module: Mapping[str, set[str]]
+  v2_generated_imports_by_module: Mapping[str, set[str]]
+  docs_by_module: Mapping[str, str]
+
+
+def get_module(dir_path: str, relative_to_dir: str) -> str:
+  """Get module that corresponds to path relative to relative_to_dir.
+
+  Args:
+    dir_path: Path to directory.
+    relative_to_dir: Get module relative to this directory.
+
+  Returns:
+    Name of module that corresponds to the given directory.
+  """
+  dir_path = dir_path[len(relative_to_dir) :]
+  # Convert path separators to '/' for easier parsing below.
+  dir_path = dir_path.replace(os.sep, '/')
+  return dir_path.replace('/', '.').strip('.')
+
+
+def generate_proxy_api_files(
+    output_files: list[str], proxy_module_root: str, output_dir: str
+):
+  """Creates __init__.py files in proxy format for the Python API.
+
+  Args:
+    output_files: List of __init__.py file paths to create.
+    proxy_module_root: Module root for proxy-import format. If specified, proxy
+      files with content like `from proxy_module_root.proxy_module import *`
+      will be created to enable import resolution under TensorFlow.
+    output_dir: output API root directory.
+  """
+  for file in output_files:
+    file_dir = os.path.dirname(file)
+    if not os.path.isdir(file_dir):
+      os.makedirs(file_dir)
+    module = get_module(file_dir, output_dir)
+    content = f'from {proxy_module_root}.{module} import *'
+    with open(file, 'w') as f:
+      f.write(content)
+
+
+def _should_skip_file(
+    file: str,
+    file_prefixes_to_strip: Sequence[str],
+    packages_to_ignore: Sequence[str],
+    module_prefix: str,
+) -> bool:
+  import_path = _get_import_path(file, file_prefixes_to_strip, module_prefix)
+  return any(import_path.startswith(package) for package in packages_to_ignore)
+
+
+def get_public_api(
+    api_mapping_files: Sequence[str],
+    file_prefixes_to_strip: Sequence[str],
+    packages_to_ignore: Sequence[str],
+    output_package: str,
+    module_prefix: str,
+) -> PublicAPI:
+  """Generates the structure of the public API from the given files.
+
+  Args:
+    api_mapping_files: List of files containing the exported API mappings and
+      docstrings.
+    file_prefixes_to_strip: A list of prefixes to strip from files when
+      determining the packages to ignore.
+    packages_to_ignore: A list of python packages that should be ignored when
+      searching for tf_exports.
+    output_package: The package to use for the imports.
+    module_prefix: A prefix to add to the non-generated imports.
+
+  Raises:
+    DocExportedTwiceError: Two docstrings are registered for the same module.
+
+  Returns:
+    The public API structure.
+  """
+  ea = exported_api.ExportedApi()
+  for f in api_mapping_files:
+    ea.read(f)
+
+  v1_entrypoints_by_module = collections.defaultdict(set)
+  v2_entrypoints_by_module = collections.defaultdict(set)
+
+  def add_exported_symbols(
+      api_names: list[str],
+      s: exported_api.ExportedSymbol,
+      entrypoints_by_module: Mapping[str, set[_Entrypoint]],
+  ):
+    for api_name in api_names:
+      index_of_last_dot = api_name.rfind('.')
+      index_of_first_dot = api_name.find('.')
+      module = output_package
+      if index_of_first_dot + 1 < index_of_last_dot:
+        module += f'.{api_name[index_of_first_dot + 1:index_of_last_dot]}'
+      name = api_name[index_of_last_dot + 1 :]
+      entrypoints_by_module[module].add(_Entrypoint(module, name, s))
+
+  for s in ea.symbols:
+    if _should_skip_file(
+        s.file_name, file_prefixes_to_strip, packages_to_ignore, module_prefix
+    ):
+      continue
+    add_exported_symbols(s.v1_apis, s, v1_entrypoints_by_module)
+    add_exported_symbols(s.v2_apis, s, v2_entrypoints_by_module)
+
+  v1_generated_imports_by_module = collections.defaultdict(set)
+  v2_generated_imports_by_module = collections.defaultdict(set)
+
+  def add_generated_imports(
+      entrypoints_by_module: Mapping[str, set[_Entrypoint]],
+      generated_imports_by_module: Mapping[str, set[str]],
+  ):
+    for module in entrypoints_by_module:
+      i = module.rfind('.')
+      if i == -1:
+        continue
+      while i != -1:
+        parent = module[:i]
+        generated_imports_by_module[parent].add(module)
+        module = parent
+        i = module.rfind('.')
+
+  add_generated_imports(
+      v1_entrypoints_by_module, v1_generated_imports_by_module
+  )
+  add_generated_imports(
+      v2_entrypoints_by_module, v2_generated_imports_by_module
+  )
+
+  docs_by_module = {}
+
+  for d in ea.docs:
+    for m in d.modules:
+      if m in docs_by_module:
+        raise DocExportedTwiceError(
+            f'Docstring at {d.file_name}:{d.line_no} is registered for {m},'
+            ' which already has a registered docstring.'
+        )
+      docs_by_module[m] = d.docstring
+
+  return PublicAPI(
+      v1_entrypoints_by_module=v1_entrypoints_by_module,
+      v2_entrypoints_by_module=v2_entrypoints_by_module,
+      v1_generated_imports_by_module=v1_generated_imports_by_module,
+      v2_generated_imports_by_module=v2_generated_imports_by_module,
+      docs_by_module=docs_by_module,
+  )
+
+
+def _get_module_docstring(
+    docs_by_module: Mapping[str, str], module: str
+) -> str:
+  if module in docs_by_module:
+    return docs_by_module[module]
+  module = module.replace('tensorflow', 'tf')
+  return f'Public API for {module} namespace'
+
+
+def _get_imports_for_module(
+    module: str,
+    output_package: str,
+    symbols_by_module: Mapping[str, set[_Entrypoint]],
+    generated_imports_by_module: Mapping[str, set[str]],
+    file_prefixes_to_strip: Sequence[str],
+    module_prefix: str,
+    use_lazy_loading: bool,
+    subpackage_rewrite: Optional[str],
+) -> str:
+  """Returns the imports for a module.
+
+  Args:
+    module: The module to get imports for.
+    output_package: The package to use for the imports.
+    symbols_by_module: The symbols that should be exposed by each module.
+    generated_imports_by_module: The sub-modules that should be exposed by each
+      module.
+    file_prefixes_to_strip: The prefixes to strip from the file names of the
+      imports.
+    module_prefix: A prefix to add to the non-generated imports.
+    use_lazy_loading: Whether to use lazy loading or not.
+    subpackage_rewrite: The subpackage to use for the imports.
+  """
+  content = ''
+  symbol_imports = list(symbols_by_module[module])
+  symbol_imports = sorted(
+      symbol_imports, key=lambda s: f'{s.exported_symbol.file_name}:{s.name}'
+  )
+  generated_imports = sorted(generated_imports_by_module[module])
+  for imp in generated_imports:
+    if subpackage_rewrite:
+      imp = imp.replace(output_package, subpackage_rewrite)
+    last_dot = imp.rfind('.')
+    if use_lazy_loading:
+      content += f"  '{imp[last_dot+1:]}': ('', '{imp}'),\n"
+    else:
+      content += f'from {imp[:last_dot]} import {imp[last_dot+1:]}\n'
+  for s in symbol_imports:
+    content += (
+        f'{s.get_import(file_prefixes_to_strip, module_prefix, use_lazy_loading=use_lazy_loading)}\n'
+    )
+  return content
+
+
+def gen_public_api(
+    output_dir: str,
+    output_package: str,
+    root_init_template: str,
+    api_version: int,
+    compat_api_versions: Sequence[int],
+    compat_init_templates: Sequence[str],
+    use_lazy_loading: bool,
+    file_prefixes_to_strip: Sequence[str],
+    mapping_files: Sequence[str],
+    packages_to_ignore: Sequence[str],
+    module_prefix: str,
+    root_file_name: str,
+    output_files: Set[str],
+):
+  """Generates the public API for tensorflow.
+
+  Args:
+    output_dir: The directory to output the files to.
+    output_package: The package to use for the imports.
+    root_init_template: The template for the root init file.
+    api_version: The version of the API to generate.
+    compat_api_versions: The versions of the compat APIs to generate.
+    compat_init_templates: The templates for the compat init files.
+    use_lazy_loading: Whether to use lazy loading or not.
+    file_prefixes_to_strip: The prefixes to strip from the file names of the
+      imports.
+    mapping_files: The mapping files created by the API Extractor.
+    packages_to_ignore: A list of python packages that should be ignored when
+      searching for tf_exports.
+    module_prefix: A prefix to add to the non-generated imports.
+    root_file_name: The file name that should be generated for the top level
+      API.
+    output_files: List of files expected to generate.
+  """
+  public_api = get_public_api(
+      mapping_files,
+      file_prefixes_to_strip,
+      packages_to_ignore,
+      output_package,
+      module_prefix,
+  )
+
+  root_entrypoints_by_module = public_api.v2_entrypoints_by_module
+  root_generated_imports_by_module = public_api.v2_generated_imports_by_module
+  if api_version == 1:
+    root_entrypoints_by_module = public_api.v1_entrypoints_by_module
+    root_generated_imports_by_module = public_api.v1_generated_imports_by_module
+
+  for compat_version in compat_api_versions:
+    compat_package = f'{output_package}.compat'
+    compat_version_package = f'{compat_package}.v{compat_version}'
+    public_api.v2_generated_imports_by_module[compat_package].add(
+        compat_version_package
+    )
+    public_api.v1_generated_imports_by_module[compat_package].add(
+        compat_version_package
+    )
+
+  _gen_init_files(
+      output_dir,
+      output_package,
+      api_version,
+      root_entrypoints_by_module,
+      root_generated_imports_by_module,
+      public_api.docs_by_module,
+      root_init_template,
+      file_prefixes_to_strip,
+      use_lazy_loading,
+      module_prefix,
+      output_files,
+      root_file_name=root_file_name,
+  )
+
+  for compat_index, compat_version in enumerate(compat_api_versions):
+    compat_output_dir = os.path.join(output_dir, 'compat', f'v{compat_version}')
+    os.makedirs(compat_output_dir, exist_ok=True)
+    compat_version = int(compat_version)
+
+    compat_entrypoints_by_module = public_api.v2_entrypoints_by_module
+    compat_generated_imports_by_module = (
+        public_api.v2_generated_imports_by_module
+    )
+    if compat_version == 1:
+      compat_entrypoints_by_module = public_api.v1_entrypoints_by_module
+      compat_generated_imports_by_module = (
+          public_api.v1_generated_imports_by_module
+      )
+
+    _gen_init_files(
+        compat_output_dir,
+        output_package,
+        compat_version,
+        compat_entrypoints_by_module,
+        compat_generated_imports_by_module,
+        public_api.docs_by_module,
+        compat_init_templates[compat_index] if compat_init_templates else '',
+        file_prefixes_to_strip,
+        use_lazy_loading,
+        module_prefix,
+        output_files,
+        subpackage_rewrite=f'{output_package}.compat.v{compat_version}',
+    )
+
+    for nested_compat_index, nested_compat_version in enumerate(
+        compat_api_versions
+    ):
+      nested_compat_version = int(nested_compat_version)
+      nested_compat_output_dir = os.path.join(
+          compat_output_dir, 'compat', f'v{nested_compat_version}'
+      )
+      nested_compat_entrypoints_by_module = public_api.v2_entrypoints_by_module
+      nested_compat_generated_imports_by_module = (
+          public_api.v2_generated_imports_by_module
+      )
+      if nested_compat_version == 1:
+        nested_compat_entrypoints_by_module = (
+            public_api.v1_entrypoints_by_module
+        )
+        nested_compat_generated_imports_by_module = (
+            public_api.v1_generated_imports_by_module
+        )
+      os.makedirs(nested_compat_output_dir, exist_ok=True)
+      gen_nested_compat_files(
+          nested_compat_output_dir,
+          output_package,
+          nested_compat_version,
+          nested_compat_entrypoints_by_module,
+          nested_compat_generated_imports_by_module,
+          public_api.docs_by_module,
+          compat_init_templates[nested_compat_index]
+          if compat_init_templates
+          else '',
+          file_prefixes_to_strip,
+          use_lazy_loading,
+          compat_api_versions,
+          module_prefix,
+          output_files,
+      )
+
+
+def _get_module_wrapper(
+    module: str,
+    output_dir: str,
+    output_package: str,
+    api_version: int,
+    symbols_by_module: Mapping[str, set[_Entrypoint]],
+    use_lazy_loading: bool,
+) -> str:
+  """Returns the module wrapper for the given module."""
+  if api_version != 1 and not use_lazy_loading:
+    return ''
+  deprecated = 'False'
+  has_lite = 'False'
+  public_apis_name = 'None'
+  if api_version == 1 and not output_dir.strip('/').endswith('compat/v1'):
+    deprecated = 'True'
+  if 'lite' in symbols_by_module and use_lazy_loading:
+    has_lite = 'True'
+  if use_lazy_loading:
+    public_apis_name = '_PUBLIC_APIS'
+  return _DEPRECATION_FOOTER % (
+      module.removeprefix(output_package).strip('.'),
+      public_apis_name,
+      deprecated,
+      has_lite,
+  )
+
+
+def _gen_init_files(
+    output_dir: str,
+    output_package: str,
+    api_version: int,
+    symbols_by_module: Mapping[str, set[_Entrypoint]],
+    generated_imports_by_module: Mapping[str, set[str]],
+    docs_by_module: Mapping[str, str],
+    root_template_path: str,
+    file_prefixes_to_strip: Sequence[str],
+    use_lazy_loading: bool,
+    module_prefix: str,
+    output_files: Set[str],
+    subpackage_rewrite: Optional[str] = None,
+    root_file_name='__init__.py',
+):
+  """Generates the __init__.py files for the given API version."""
+  modules = set(symbols_by_module.keys())
+  modules.update(generated_imports_by_module.keys())
+  for module in modules:
+    if len(module) < len(output_package):
+      continue
+    module_relative_to_package = module[len(output_package) + 1 :]
+    module_path = os.path.join(
+        output_dir, module_relative_to_package.replace('.', '/')
+    )
+    os.makedirs(module_path, exist_ok=True)
+    module_file_path = os.path.join(
+        module_path,
+        root_file_name if not module_relative_to_package else '__init__.py',
+    )
+    module_file_path = os.path.normpath(module_file_path)
+    if module_file_path not in output_files:
+      raise AssertionError(
+          f'Exported api attempted to write to "{module_file_path}" but it is'
+          ' not in output_files.'
+      )
+    with open(module_file_path, 'w') as f:
+      module_imports = _get_imports_for_module(
+          module,
+          output_package,
+          symbols_by_module,
+          generated_imports_by_module,
+          file_prefixes_to_strip,
+          module_prefix,
+          use_lazy_loading,
+          subpackage_rewrite,
+      )
+      if use_lazy_loading:
+        module_imports = _LAZY_LOADING_MODULE_TEXT_TEMPLATE % module_imports
+      # If this module is the root and there is a root template, use it
+      if module == output_package and root_template_path:
+        with open(root_template_path, 'r') as template:
+          content = template.read()
+          content = content.replace('# API IMPORTS PLACEHOLDER', module_imports)
+
+          underscore_elements = [
+              s.name
+              for s in symbols_by_module[module]
+              if s.name.startswith('_')
+          ]
+          for i in generated_imports_by_module[module]:
+            module_name = i[i.rfind('.') + 1 :]
+            if module_name.startswith('_'):
+              underscore_elements.append(module_name)
+
+          root_module_footer = f"""
+_names_with_underscore = [{', '.join(sorted([f"'{s}'" for s in underscore_elements]))}]
+__all__ = [_s for _s in dir() if not _s.startswith('_')]
+__all__.extend([_s for _s in _names_with_underscore])
+          """
+
+          content = content.replace('# __all__ PLACEHOLDER', root_module_footer)
+
+          content = content.replace(
+              '# WRAPPER_PLACEHOLDER',
+              _get_module_wrapper(
+                  module,
+                  output_dir,
+                  output_package,
+                  api_version,
+                  symbols_by_module,
+                  use_lazy_loading,
+              ),
+          )
+
+          f.write(content)
+          continue
+
+      f.write(
+          _GENERATED_FILE_HEADER % _get_module_docstring(docs_by_module, module)
+      )
+
+      f.write(module_imports)
+
+      f.write(
+          _get_module_wrapper(
+              module,
+              output_dir,
+              output_package,
+              api_version,
+              symbols_by_module,
+              use_lazy_loading,
+          )
+      )
+
+
+def gen_nested_compat_files(
+    output_dir: str,
+    output_package: str,
+    api_version: int,
+    symbols_by_module: Mapping[str, set[_Entrypoint]],
+    generated_imports_by_module: Mapping[str, set[str]],
+    docs_by_module: Mapping[str, str],
+    root_template_path: str,
+    file_prefixes_to_strip: Sequence[str],
+    use_lazy_loading: bool,
+    compat_versions: Sequence[int],
+    module_prefix: str,
+    output_files: Set[str],
+):
+  """Generates the nested compat __init__.py files."""
+  nested_compat_symbols_by_module: dict[str, set[_Entrypoint]] = {}
+  nested_generated_imports_by_module: dict[str, set[str]] = {}
+  compat_module = f'{output_package}.compat'
+  # The nested compat files should only generate imports for the nested root
+  # package, and its corresponding compat package.
+  if output_package in symbols_by_module:
+    nested_compat_symbols_by_module[output_package] = symbols_by_module[
+        output_package
+    ]
+  if compat_module in symbols_by_module:
+    nested_compat_symbols_by_module[compat_module] = symbols_by_module[
+        compat_module
+    ]
+  if output_package in generated_imports_by_module:
+    nested_generated_imports_by_module[output_package] = (
+        generated_imports_by_module[output_package]
+    )
+  if compat_module in generated_imports_by_module:
+    nested_generated_imports_by_module[compat_module] = (
+        generated_imports_by_module[compat_module]
+    )
+
+  _gen_init_files(
+      output_dir,
+      output_package,
+      api_version,
+      nested_compat_symbols_by_module,
+      nested_generated_imports_by_module,
+      docs_by_module,
+      root_template_path,
+      file_prefixes_to_strip,
+      use_lazy_loading,
+      module_prefix,
+      output_files,
+      f'{compat_module}.v{api_version}',
+  )
+
+  for compat_version in compat_versions:
+    nested_generated_imports_by_module[compat_module].add(
+        f'{output_package}.compat.v{compat_version}'
+    )
+
+
+def main(argv: Sequence[str]) -> None:
+  if not _OUTPUT_DIR.value or not _OUTPUT_FILES.value:
+    raise app.UsageError('--output_dir and --output_files are required')
+
+  if _PROXY_MODULE_ROOT.value:
+    generate_proxy_api_files(
+        _OUTPUT_FILES.value, _PROXY_MODULE_ROOT.value, _OUTPUT_DIR.value
+    )
+    return
+
+  output_files = [os.path.normpath(f) for f in _OUTPUT_FILES.value]
+
+  for out_file in output_files:
+    with open(out_file, 'w') as f:
+      f.write('')
+
+  gen_public_api(
+      _OUTPUT_DIR.value,
+      _OUTPUT_PACKAGE.value,
+      _ROOT_INIT_TEMPLATE.value,
+      _API_VERSION.value,
+      [int(v) for v in _COMPAT_API_VERSIONS.value],
+      _COMPAT_INIT_TEMPLATES.value,
+      _USE_LAZY_LOADING.value,
+      _FILE_PREFIXES_TO_STRIP.value,
+      argv[1:],
+      _PACKAGES_TO_IGNORE.value,
+      _MODULE_PREFIX.value,
+      _ROOT_FILE_PATH.value,
+      set(output_files),
+  )

tensorflow/python/tools/api/generator2/generator/generator_test.py

@@ -0,0 +1,481 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+import collections
+import os
+
+from absl.testing import absltest
+from absl.testing import parameterized
+
+from tensorflow.python.tools.api.generator2.generator import generator
+from tensorflow.python.tools.api.generator2.shared import exported_api
+
+tensor_es = exported_api.ExportedSymbol(
+    file_name='tf/python/framework/tensor.py',
+    line_no=1,
+    symbol_name='Tensor',
+    v1_apis=('tf.Tensor',),
+    v2_apis=(
+        'tf.Tensor',
+        'tf.experimental.numpy.ndarray',
+    ),
+)
+
+test_data = {
+    'tf/python/framework/tensor_mapping.json': exported_api.ExportedApi(
+        docs=[],
+        symbols=[tensor_es],
+    ),
+    'tf/python/framework/test_ops.json': exported_api.ExportedApi(
+        docs=[],
+        symbols=[
+            exported_api.ExportedSymbol(
+                file_name='tf/python/framework/test_ops.py',
+                line_no=2,
+                symbol_name='a',
+                v1_apis=(),
+                v2_apis=('a',),
+            )
+        ],
+    ),
+}
+
+
+def write_test_data(tmp_dir: str):
+  for f in test_data:
+    file_name = os.path.join(tmp_dir, f)
+    os.makedirs(os.path.dirname(file_name), exist_ok=True)
+    test_data[f].write(file_name)
+
+
+class GeneratorTest(parameterized.TestCase):
+
+  def test_get_public_api(self):
+    tmp_dir = self.create_tempdir()
+    write_test_data(tmp_dir.full_path)
+
+    expected_tensor_top_level = generator._Entrypoint(
+        module='tf',
+        name='Tensor',
+        exported_symbol=tensor_es,
+    )
+
+    expected = generator.PublicAPI(
+        v1_entrypoints_by_module=collections.defaultdict(set),
+        v2_entrypoints_by_module=collections.defaultdict(set),
+        v1_generated_imports_by_module=collections.defaultdict(set),
+        v2_generated_imports_by_module=collections.defaultdict(set),
+        docs_by_module={},
+    )
+    expected.v1_entrypoints_by_module['tf'].add(expected_tensor_top_level)
+    expected.v2_entrypoints_by_module['tf'].add(expected_tensor_top_level)
+    expected.v2_entrypoints_by_module['tf.experimental.numpy'].add(
+        generator._Entrypoint(
+            module='tf.experimental.numpy',
+            name='ndarray',
+            exported_symbol=tensor_es,
+        )
+    )
+    expected.v2_generated_imports_by_module['tf'].add('tf.experimental')
+    expected.v2_generated_imports_by_module['tf.experimental'].add(
+        'tf.experimental.numpy'
+    )
+
+    got = generator.get_public_api(
+        [os.path.join(tmp_dir, f) for f in test_data],
+        file_prefixes_to_strip=[tmp_dir.full_path],
+        packages_to_ignore=['tf.python.framework.test_ops'],
+        output_package='tf',
+        module_prefix=''
+    )
+
+    self.assertEqual(
+        expected,
+        got,
+    )
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name='normal_file',
+          entrypoint=generator._Entrypoint(
+              module='tf.io',
+              name='decode_csv',
+              exported_symbol=exported_api.ExportedSymbol(
+                  file_name='tf/python/ops/parsing_ops.py',
+                  line_no=10,
+                  symbol_name='decode_csv_v2',
+                  v1_apis=[],
+                  v2_apis=['tf.io.decode_csv'],
+              ),
+          ),
+          prefixes_to_strip=[],
+          expected='tf.python.ops.parsing_ops',
+      ),
+      dict(
+          testcase_name='genfile',
+          entrypoint=generator._Entrypoint(
+              module='tf.io',
+              name='decode_proto_v2',
+              exported_symbol=exported_api.ExportedSymbol(
+                  file_name=(
+                      'bazel-out/genfiles/tf/python/ops/gen_decode_proto_ops.py'
+                  ),
+                  line_no=20,
+                  symbol_name='decode_proto_v2',
+                  v1_apis=[],
+                  v2_apis=['tf.io.decode_proto_v2'],
+              ),
+          ),
+          prefixes_to_strip=['bazel-out/genfiles'],
+          expected='tf.python.ops.gen_decode_proto_ops',
+      ),
+  )
+  def test_get_import_path(self, entrypoint, prefixes_to_strip, expected):
+    self.assertEqual(
+        expected,
+        generator._get_import_path(
+            entrypoint.exported_symbol.file_name, prefixes_to_strip, ''
+        ),
+    )
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name='direct',
+          entrypoint=generator._Entrypoint(
+              module='tf',
+              name='Tensor',
+              exported_symbol=tensor_es,
+          ),
+          use_lazy_loading=False,
+          expected='from tf.python.framework.tensor import Tensor # line: 1',
+      ),
+      dict(
+          testcase_name='alias',
+          entrypoint=generator._Entrypoint(
+              module='tf.io',
+              name='decode_csv',
+              exported_symbol=exported_api.ExportedSymbol(
+                  file_name='tf/python/ops/parsing_ops.py',
+                  line_no=10,
+                  symbol_name='decode_csv_v2',
+                  v1_apis=[],
+                  v2_apis=['tf.io.decode_csv'],
+              ),
+          ),
+          use_lazy_loading=False,
+          expected=(
+              'from tf.python.ops.parsing_ops import decode_csv_v2 as'
+              ' decode_csv # line: 10'
+          ),
+      ),
+      dict(
+          testcase_name='direct_lazy',
+          entrypoint=generator._Entrypoint(
+              module='tf',
+              name='Tensor',
+              exported_symbol=tensor_es,
+          ),
+          use_lazy_loading=True,
+          expected=(
+              "  'Tensor': ('tf.python.framework.tensor', 'Tensor'), # line: 1"
+          ),
+      ),
+      dict(
+          testcase_name='alias_lazy',
+          entrypoint=generator._Entrypoint(
+              module='tf.io',
+              name='decode_csv',
+              exported_symbol=exported_api.ExportedSymbol(
+                  file_name='tf/python/ops/parsing_ops.py',
+                  line_no=10,
+                  symbol_name='decode_csv_v2',
+                  v1_apis=[],
+                  v2_apis=['tf.io.decode_csv'],
+              ),
+          ),
+          use_lazy_loading=True,
+          expected=(
+              "  'decode_csv': ('tf.python.ops.parsing_ops',"
+              " 'decode_csv_v2'), # line: 10"
+          ),
+      ),
+  )
+  def test_entrypoint_get_import(self, entrypoint, use_lazy_loading, expected):
+    self.assertEqual(expected, entrypoint.get_import([], '', use_lazy_loading))
+
+  def test_get_module(self):
+    self.assertEqual(
+        'keras.losses',
+        generator.get_module(
+            'bazel/tensorflow/keras/losses/', 'bazel/tensorflow'
+        ),
+    )
+
+  def test_generate_proxy_api_files(self):
+    tmp_dir = self.create_tempdir()
+    proxy_file = os.path.join(tmp_dir, 'tensorflow/keras/losses/__init__.py')
+    generator.generate_proxy_api_files(
+        [proxy_file], 'keras', os.path.join(tmp_dir, 'tensorflow/keras')
+    )
+    self.assertTrue(os.path.isfile(proxy_file))
+    with open(proxy_file, 'r') as f:
+      self.assertEqual('from keras.losses import *', f.read())
+
+  def test_get_module_docstring(self):
+    docs_by_module = {
+        'io': 'io docs',
+    }
+    self.assertEqual(
+        'io docs', generator._get_module_docstring(docs_by_module, 'io')
+    )
+    self.assertEqual(
+        'Public API for math namespace',
+        generator._get_module_docstring(docs_by_module, 'math'),
+    )
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name='static_imports',
+          use_lazy_loading=False,
+          subpackage_rewrite=None,
+          expected="""from tf import io
+from tf.python.framework.tensor import Tensor # line: 1
+""",
+      ),
+      dict(
+          testcase_name='lazy_imports',
+          use_lazy_loading=True,
+          subpackage_rewrite=None,
+          expected="""  'io': ('', 'tf.io'),
+  'Tensor': ('tf.python.framework.tensor', 'Tensor'), # line: 1
+""",
+      ),
+      dict(
+          testcase_name='subpackage_rewrite',
+          use_lazy_loading=False,
+          subpackage_rewrite='tf.compat.v1',
+          expected="""from tf.compat.v1 import io
+from tf.python.framework.tensor import Tensor # line: 1
+""",
+      ),
+  )
+  def test_get_imports_for_module(
+      self, use_lazy_loading, subpackage_rewrite, expected
+  ):
+    symbols_by_module = {
+        'tf': {
+            generator._Entrypoint(
+                module='tf', name='Tensor', exported_symbol=tensor_es
+            )
+        }
+    }
+    generated_imports_by_module = {'tf': {'tf.io'}}
+    self.assertEqual(
+        expected,
+        generator._get_imports_for_module(
+            'tf',
+            'tf',
+            symbols_by_module,
+            generated_imports_by_module,
+            [],
+            '',
+            use_lazy_loading,
+            subpackage_rewrite,
+        ),
+    )
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name='empty_prefixes_and_packages',
+          file='tf/python/framework/test_ops.py',
+          file_prefixes_to_strip=[],
+          packages_to_ignore=[],
+          should_skip=False,
+      ),
+      dict(
+          testcase_name='empty_prefix_nonempty_package',
+          file='tf/python/framework/test_ops.py',
+          file_prefixes_to_strip=[],
+          packages_to_ignore=['tf.python.framework.test_ops'],
+          should_skip=True,
+      ),
+      dict(
+          testcase_name='nonempty_prefix_empty_package',
+          file='gen/tf/python/framework/test_ops.py',
+          file_prefixes_to_strip=['gen/'],
+          packages_to_ignore=[],
+          should_skip=False,
+      ),
+      dict(
+          testcase_name='nonempty_prefix_nonempty_package',
+          file='gen/tf/python/framework/test_ops.py',
+          file_prefixes_to_strip=['gen/'],
+          packages_to_ignore=['tf.python.framework.test_ops'],
+          should_skip=True,
+      ),
+      dict(
+          testcase_name='non_matching_prefix_and_package',
+          file='tf/python/ops/test_ops.py',
+          file_prefixes_to_strip=['gen/'],
+          packages_to_ignore=['tf.python.framework.test_ops'],
+          should_skip=False,
+      ),
+  )
+  def test_should_skip_file(
+      self, file, file_prefixes_to_strip, packages_to_ignore, should_skip
+  ):
+    self.assertEqual(
+        should_skip,
+        generator._should_skip_file(
+            file, file_prefixes_to_strip, packages_to_ignore, '',
+        ),
+    )
+
+  @parameterized.named_parameters(
+      dict(
+          testcase_name='default',
+          root_file_name=None,
+      ),
+      dict(
+          testcase_name='renamed_root',
+          root_file_name='v2.py',
+      ),
+  )
+  def test_gen_init_files(self, root_file_name):
+    output_dir = self.create_tempdir()
+    mapping_dir = self.create_tempdir()
+    write_test_data(mapping_dir.full_path)
+
+    file_prefixes_to_strip = [mapping_dir.full_path]
+
+    public_api = generator.get_public_api(
+        [os.path.join(mapping_dir, f) for f in test_data],
+        file_prefixes_to_strip=file_prefixes_to_strip,
+        packages_to_ignore=['tf.python.framework.test_ops'],
+        output_package='tf',
+        module_prefix='',
+    )
+
+    paths_expected = [
+        root_file_name if root_file_name else '__init__.py',
+        'experimental/__init__.py',
+        'experimental/numpy/__init__.py',
+    ]
+    paths_expected = set(
+        [
+            os.path.normpath(os.path.join(output_dir, path))
+            for path in paths_expected
+        ]
+    )
+    if root_file_name is None:
+      generator._gen_init_files(
+          output_dir,
+          'tf',
+          2,
+          public_api.v2_entrypoints_by_module,
+          public_api.v2_generated_imports_by_module,
+          public_api.docs_by_module,
+          '',
+          file_prefixes_to_strip,
+          False,
+          '',
+          paths_expected,
+      )
+    else:
+      generator._gen_init_files(
+          output_dir,
+          'tf',
+          2,
+          public_api.v2_entrypoints_by_module,
+          public_api.v2_generated_imports_by_module,
+          public_api.docs_by_module,
+          '',
+          file_prefixes_to_strip,
+          False,
+          '',
+          paths_expected,
+          root_file_name=root_file_name,
+      )
+    expected_init_path = os.path.join(
+        output_dir.full_path,
+        root_file_name if root_file_name else '__init__.py',
+    )
+    self.assertTrue(os.path.exists(expected_init_path))
+    with open(expected_init_path, 'r') as f:
+      self.assertEqual(
+          f.read(),
+          """# This file is MACHINE GENERATED! Do not edit.
+# Generated by: tensorflow/python/tools/api/generator2/generator/generator.py script.
+\"""Public API for tf namespace
+\"""
+
+import sys as _sys
+
+from tf import experimental
+from tf.python.framework.tensor import Tensor # line: 1
+""",
+      )
+    expected_numpy_path = os.path.join(
+        output_dir.full_path, 'experimental/numpy/__init__.py'
+    )
+    self.assertTrue(os.path.exists(expected_numpy_path))
+    with open(expected_numpy_path, 'r') as f:
+      self.assertEqual(
+          f.read(),
+          """# This file is MACHINE GENERATED! Do not edit.
+# Generated by: tensorflow/python/tools/api/generator2/generator/generator.py script.
+\"""Public API for tf.experimental.numpy namespace
+\"""
+
+import sys as _sys
+
+from tf.python.framework.tensor import Tensor as ndarray # line: 1
+""",
+      )
+
+  def testRaisesOnNotExpectedFile(self):
+    output_dir = self.create_tempdir()
+    mapping_dir = self.create_tempdir()
+    write_test_data(mapping_dir.full_path)
+
+    file_prefixes_to_strip = [mapping_dir.full_path]
+
+    public_api = generator.get_public_api(
+        [os.path.normpath(os.path.join(mapping_dir, f)) for f in test_data],
+        file_prefixes_to_strip=file_prefixes_to_strip,
+        packages_to_ignore=['tf.python.framework.test_ops'],
+        output_package='tf',
+        module_prefix='',
+    )
+
+    with self.assertRaisesRegex(
+        AssertionError, 'Exported api attempted to write to'
+    ):
+      generator._gen_init_files(
+          output_dir,
+          'tf',
+          2,
+          public_api.v2_entrypoints_by_module,
+          public_api.v2_generated_imports_by_module,
+          public_api.docs_by_module,
+          '',
+          file_prefixes_to_strip,
+          False,
+          '',
+          [],
+      )
+
+
+if __name__ == '__main__':
+  absltest.main()

tensorflow/python/tools/api/generator2/generator/main.py

@@ -0,0 +1,22 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+"""Binary for generating the API for tensorflow."""
+
+from absl import app
+
+from tensorflow.python.tools.api.generator2.generator import generator
+
+if __name__ == '__main__':
+  app.run(generator.main)

tensorflow/python/tools/api/generator2/patterns.bzl

@@ -0,0 +1,66 @@
+"""Support for working with patterns and matching."""
+
+def Pattern(pattern):
+    """Compiles pattern into a Pattern struct.
+
+    Args:
+        pattern: Bazel Target pattern
+
+    Returns:
+        Pattern struct
+    """
+    if pattern.endswith("/..."):
+        return struct(
+            label = Label(pattern.removesuffix("/...")),
+            subpackages = True,
+        )
+    return struct(
+        label = Label(pattern),
+        subpackages = False,
+    )
+
+def compile_patterns(patterns):
+    """Compiles each string into a Pattern struct.
+
+    Args:
+        patterns: Iterable of Bazel Target pattern strings
+
+    Returns:
+        List of Pattern structs
+    """
+    return [Pattern(pattern) for pattern in patterns]
+
+def matches(pattern, label):
+    """Checks if patterns includes label.
+
+    Args:
+        pattern: A Pattern struct
+        label: Bazel Label object
+
+    Returns:
+        True if pattern includes label, False otherwise
+    """
+    if pattern.label.workspace_name != label.workspace_name:
+        return False
+    if pattern.subpackages:
+        return label.package == pattern.label.package or label.package.startswith(pattern.label.package + "/")
+    if pattern.label.package != label.package:
+        return False
+    if pattern.label.name == "all" or pattern.label.name == "*":
+        return True
+    return pattern.label.name == label.name
+
+def any_match(patterns, label):
+    """Whether any Pattern in patterns include labels.
+
+    Args:
+        patterns: An iterable of Pattern structs
+        label: Bazel Label object
+
+    Returns:
+        True if any pattern includes label, False otherwise
+    """
+    for pattern in patterns:
+        if matches(pattern, label):
+            return True
+    return False

tensorflow/python/tools/api/generator2/shared/BUILD

@@ -0,0 +1,25 @@
+load("//tensorflow:pytype.default.bzl", "pytype_strict_library")
+load("//tensorflow:tensorflow.bzl", "py_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//tensorflow/python/tools/api/generator2:__subpackages__"],
+    licenses = ["notice"],
+)
+
+pytype_strict_library(
+    name = "exported_api",
+    srcs = ["exported_api.py"],
+)
+
+py_test(
+    name = "exported_api_test",
+    srcs = ["exported_api_test.py"],
+    strict_deps = True,
+    tags = ["no_pip"],
+    deps = [
+        ":exported_api",
+        # copybara:uncomment "//third_party/py/google/protobuf:use_fast_cpp_protos",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)

tensorflow/python/tools/api/generator2/shared/exported_api.py

@@ -0,0 +1,113 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Reads and writes files with TF Python exports metadata."""
+
+from collections.abc import Iterable, Sequence
+import json
+from typing import Any, NamedTuple
+
+
+class ExportedSymbol(NamedTuple):
+  """Information about a single tf_export instance."""
+
+  file_name: str
+  line_no: int
+  symbol_name: str
+  v1_apis: tuple[str, ...]
+  v2_apis: tuple[str, ...]
+
+  @classmethod
+  def create(
+      cls, *, v1_apis: Sequence[str], v2_apis: Sequence[str], **kwargs
+  ) -> "ExportedSymbol":
+    return cls(v1_apis=tuple(v1_apis), v2_apis=tuple(v2_apis), **kwargs)
+
+
+class ExportedDoc(NamedTuple):
+  """Information about an export Module docstring."""
+
+  file_name: str
+  line_no: int
+  modules: tuple[str, ...]
+  docstring: str
+
+  @classmethod
+  def create(cls, *, modules: Sequence[str], **kwargs) -> "ExportedDoc":
+    return cls(modules=tuple(modules), **kwargs)
+
+
+class ExportedApi(object):
+  """ExportedApi is a collection of ExportedSymbols."""
+
+  _docs: set[ExportedDoc]
+  _symbols: set[ExportedSymbol]
+
+  def __init__(
+      self,
+      *,
+      docs: Iterable[ExportedDoc] = (),
+      symbols: Iterable[ExportedSymbol] = (),
+  ):
+    self._docs = set(docs)
+    self._symbols = set(symbols)
+
+  def write(self, filename: str, **kwargs) -> None:
+    """Writes exports to filename."""
+    with open(filename, mode="w", encoding="utf-8") as f:
+      json.dump(
+          {
+              "docs": [d._asdict() for d in sorted(self.docs)],
+              "symbols": [s._asdict() for s in sorted(self.symbols)],
+          },
+          f,
+          **kwargs,
+      )
+
+  def read(self, filename: str) -> None:
+    """Reads exports from filename."""
+    with open(filename, mode="r", encoding="utf-8") as f:
+      data = json.load(f)
+      self._docs.update(ExportedDoc.create(**d) for d in data["docs"])
+      self._symbols.update(ExportedSymbol.create(**s) for s in data["symbols"])
+
+  def add_symbol(self, export: ExportedSymbol) -> None:
+    self._symbols.add(export)
+
+  def add_doc(self, export: ExportedDoc) -> None:
+    self._docs.add(export)
+
+  @property
+  def docs(self) -> Iterable[ExportedDoc]:
+    return self._docs
+
+  @property
+  def symbols(self) -> Iterable[ExportedSymbol]:
+    return self._symbols
+
+  def __str__(self) -> str:
+    return json.dumps({
+        "docs": [d._asdict() for d in sorted(self.docs)],
+        "symbols": [s._asdict() for s in sorted(self.symbols)],
+    })
+
+  def __repr__(self) -> str:
+    return str(self)
+
+  def __eq__(self, o: Any) -> bool:
+    return (
+        type(self) is type(o)
+        and self.docs == o.docs
+        and self.symbols == o.symbols
+    )

tensorflow/python/tools/api/generator2/shared/exported_api_test.py

@@ -0,0 +1,61 @@
+# Copyright 2023 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+from tensorflow.python.platform import test
+from tensorflow.python.tools.api.generator2.shared import exported_api
+
+_EXPORTS = exported_api.ExportedApi(
+    docs=[
+        exported_api.ExportedDoc(
+            file_name="tf/python/framework/tensor.py",
+            line_no=0,
+            modules=("tf",),
+            docstring="This is a docstring",
+        ),
+    ],
+    symbols=[
+        exported_api.ExportedSymbol(
+            file_name="tf/python/framework/tensor.py",
+            line_no=139,
+            symbol_name="Tensor",
+            v1_apis=("tf.Tensor",),
+            v2_apis=(
+                "tf.Tensor",
+                "tf.experimental.numpy.ndarray",
+            ),
+        ),
+        exported_api.ExportedSymbol(
+            file_name="tf/python/framework/tensor.py",
+            line_no=770,
+            symbol_name="Tensor",
+            v1_apis=("tf.enable_tensor_equality",),
+            v2_apis=(),
+        ),
+    ],
+)
+
+
+class ExportedApiTest(test.TestCase):
+
+  def test_read_write(self):
+    filename = self.get_temp_dir() + "/test_write.json"
+    _EXPORTS.write(filename)
+    e = exported_api.ExportedApi()
+    e.read(filename)
+
+    self.assertEqual(e, _EXPORTS)
+
+
+if __name__ == "__main__":
+  test.main()