chore: import upstream snapshot with attribution

tensorflow/python/tools/api/generator/BUILD

@@ -0,0 +1,105 @@
+# Description:
+# Scripts used to generate TensorFlow Python API.
+
+load("@bazel_skylib//:bzl_library.bzl", "bzl_library")
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:tensorflow.bzl", "py_test")
+load("//tensorflow/python/tools/api/generator:api_gen.bzl", "TENSORFLOW_API_GEN_PACKAGES")
+load("//tensorflow/python/tools/api/generator:api_init_files.bzl", "KERAS_API_INIT_FILES", "TENSORFLOW_API_INIT_FILES")
+load("//tensorflow/python/tools/api/generator:api_init_files_v1.bzl", "KERAS_API_INIT_FILES_V1", "TENSORFLOW_API_INIT_FILES_V1")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+exports_files(
+    [
+        "create_python_api.py",
+    ],
+)
+
+py_library(
+    name = "create_python_api",
+    srcs = ["create_python_api.py"],
+    strict_deps = True,
+    visibility = ["//visibility:public"],
+    deps = [
+        ":doc_srcs",
+        "//tensorflow/python/util:tf_decorator_py",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "doc_srcs",
+    srcs = ["doc_srcs.py"],
+    strict_deps = True,
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_test(
+    name = "create_python_api_test",
+    srcs = ["create_python_api_test.py"],
+    strict_deps = True,
+    deps = [
+        ":create_python_api",
+        # copybara:uncomment "//third_party/py/google/protobuf:use_fast_cpp_protos",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/util:tf_decorator_py",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_test(
+    name = "tensorflow_doc_srcs_test",
+    srcs = ["doc_srcs_test.py"],
+    args = [
+        "--package=tensorflow.python",
+        "--api_name=tensorflow",
+    ] + KERAS_API_INIT_FILES + KERAS_API_INIT_FILES_V1 + TENSORFLOW_API_INIT_FILES + TENSORFLOW_API_INIT_FILES_V1,
+    main = "doc_srcs_test.py",
+    strict_deps = True,
+    deps = [
+        ":doc_srcs",
+        # copybara:uncomment "//third_party/py/google/protobuf:use_fast_cpp_protos",
+        "//tensorflow/python:no_contrib",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_test(
+    name = "output_init_files_test",
+    srcs = ["output_init_files_test.py"],
+    args = [
+        "--packages=" + ",".join(TENSORFLOW_API_GEN_PACKAGES),
+    ],
+    data = [
+        "api_init_files.bzl",
+        "api_init_files_v1.bzl",
+    ],
+    strict_deps = True,
+    tags = [
+        "no_pip",
+    ],
+    deps = [
+        # copybara:uncomment "//third_party/py/google/protobuf:use_fast_cpp_protos",
+        "//tensorflow/dtensor/python:dtensor",
+        "//tensorflow/lite/python:analyzer",
+        "//tensorflow/lite/python:lite",
+        "//tensorflow/lite/python/authoring",
+        "//tensorflow/python:modules_with_exports",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/platform:resource_loader",
+        "//tensorflow/python/util:tf_decorator_py",
+    ],
+)
+
+bzl_library(
+    name = "api_init_files",
+    srcs = ["api_init_files.bzl"],
+    visibility = ["//tensorflow/python/tools/api/generator2:__subpackages__"],
+)

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

@@ -0,0 +1,274 @@
+"""Targets for generating TensorFlow Python API __init__.py files."""
+
+load("@xla//third_party/rules_python/python:defs.bzl", "py_binary")
+load("//tensorflow:tensorflow.bzl", "if_oss")
+load("//tensorflow:tensorflow.default.bzl", "if_indexing_source_code")
+load("//tensorflow/python/tools/api/generator:api_init_files.bzl", "TENSORFLOW_API_INIT_FILES")
+
+TENSORFLOW_API_GEN_PACKAGES = [
+    "tensorflow.python",
+    "tensorflow.compiler.mlir.quantization.tensorflow.python.quantize_model",
+    "tensorflow.compiler.mlir.quantization.tensorflow.python.representative_dataset",
+    "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.analyzer",
+    "tensorflow.lite.python.lite",
+    "tensorflow.lite.python.authoring.authoring",
+    "tensorflow.python.modules_with_exports",
+]
+
+def get_compat_files(
+        file_paths,
+        compat_api_version):
+    """Prepends compat/v<compat_api_version> to file_paths."""
+    return ["compat/v%d/%s" % (compat_api_version, f) for f in file_paths]
+
+def get_nested_compat_files(compat_api_versions):
+    """Return __init__.py file paths for files under nested compat modules.
+
+    A nested compat module contains two __init__.py files:
+      1. compat/vN/compat/vK/__init__.py
+      2. compat/vN/compat/vK/compat/__init__.py
+
+    Args:
+      compat_api_versions: list of compat versions.
+
+    Returns:
+      List of __init__.py file paths to include under nested compat modules.
+    """
+    files = []
+    for v in compat_api_versions:
+        files.extend([
+            "compat/v%d/compat/v%d/__init__.py" % (v, sv)
+            for sv in compat_api_versions
+        ])
+        files.extend([
+            "compat/v%d/compat/v%d/compat/__init__.py" % (v, sv)
+            for sv in compat_api_versions
+        ])
+    return files
+
+def gen_api_init_files(
+        name,
+        output_files = TENSORFLOW_API_INIT_FILES,
+        root_init_template = None,
+        srcs = [],
+        api_name = "tensorflow",
+        api_version = 2,
+        compat_api_versions = [],
+        compat_init_templates = [],
+        packages = TENSORFLOW_API_GEN_PACKAGES,
+        package_deps = [
+            "//tensorflow/python:no_contrib",
+            "//tensorflow/python:modules_with_exports",
+            "//tensorflow/lite/python:analyzer",
+            "//tensorflow/lite/python:lite",
+            "//tensorflow/lite/python/authoring",
+        ],
+        output_package = "tensorflow",
+        output_dir = "",
+        root_file_name = "__init__.py",
+        proxy_module_root = None,
+        api_packages_file_name = None,
+        visibility = []):
+    """Creates API directory structure and __init__.py files.
+
+    Creates a genrule that generates a directory structure with __init__.py
+    files that import all exported modules (i.e. modules with tf_export
+    decorators).
+
+    Args:
+      name: name of genrule to create.
+      output_files: 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.
+      root_init_template: Python init file that should be used as template for
+        root __init__.py file. "# API IMPORTS PLACEHOLDER" comment inside this
+        template will be replaced with root imports collected by this genrule.
+      srcs: genrule sources. If passing root_init_template, the template file
+        must be included in sources.
+      api_name: Name of the project that you want to generate API files for
+        (e.g. "tensorflow").
+      api_version: TensorFlow API version to generate. Must be either 1 or 2.
+      compat_api_versions: Older TensorFlow API versions to generate under
+        compat/ directory.
+      compat_init_templates: Python init file that should be used as template
+        for top level __init__.py files under compat/vN directories.
+        "# API IMPORTS PLACEHOLDER" comment inside this
+        template will be replaced with root imports collected by this genrule.
+      packages: Python packages containing the @tf_export decorators you want to
+        process
+      package_deps: Python library target containing your packages.
+      output_package: Package where generated API will be added to.
+      output_dir: Subdirectory to output API to.
+        If non-empty, must end with '/'.
+      root_file_name: Name of the root file with all the root imports.
+      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.
+      api_packages_file_name: Name of the file with the list of all API packages. Stores in output_dir.
+      visibility: Visibility of the rule.
+    """
+    root_init_template_flag = ""
+    if root_init_template:
+        root_init_template_flag = "--root_init_template=" + root_init_template
+
+    primary_package = packages[0]
+    api_gen_binary_target = ("create_" + primary_package + "_api_%s") % name
+    py_binary(
+        name = api_gen_binary_target,
+        srcs = ["//tensorflow/python/tools/api/generator:create_python_api.py"],
+        main = "//tensorflow/python/tools/api/generator:create_python_api.py",
+        python_version = "PY3",
+        srcs_version = "PY3",
+        visibility = ["//visibility:public"],
+        deps = package_deps + [
+            "//tensorflow/python/util:tf_decorator_py",
+            "//tensorflow/python/util:tf_export",
+            "//tensorflow/python/util:module_wrapper",
+            "//tensorflow/python/tools/api/generator:doc_srcs",
+        ],
+    )
+    if 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"]
+
+    # Replace name of root file with root_file_name.
+    output_files = [
+        root_file_name if f == "__init__.py" else f
+        for f in output_files
+    ]
+    all_output_files = ["%s%s" % (output_dir, f) for f in output_files]
+    compat_api_version_flags = ""
+    for compat_api_version in compat_api_versions:
+        compat_api_version_flags += " --compat_apiversion=%d" % compat_api_version
+
+    if api_packages_file_name:
+        api_packages_path = "%s%s" % (output_dir, api_packages_file_name)
+    else:
+        api_packages_path = None
+
+    compat_init_template_flags = ""
+    for compat_init_template in compat_init_templates:
+        compat_init_template_flags += (
+            " --compat_init_template=$(location %s)" % compat_init_template
+        )
+
+    flags = [
+        root_init_template_flag,
+        "--apidir=$(@D)/" + output_dir,
+        "--apiname=" + api_name,
+        "--apiversion=" + str(api_version),
+        compat_api_version_flags,
+        compat_init_template_flags,
+        "--packages=" + ",".join(packages),
+        "--output_package=" + output_package,
+        "--use_relative_imports=True",
+    ]
+    if proxy_module_root != None:
+        flags.append("--proxy_module_root=" + proxy_module_root)
+
+    # copybara:uncomment_begin(configurable API loading)
+    # native.vardef("TF_API_INIT_LOADING", "default")
+    # loading_value = "$(TF_API_INIT_LOADING)"
+    # copybara:uncomment_end_and_comment_begin
+    loading_value = "default"
+    # copybara:comment_end
+
+    api_gen_rule(
+        name = name,
+        outs = all_output_files,
+        srcs = srcs,
+        flags = flags,
+        api_gen_binary_target = ":" + api_gen_binary_target,
+        api_packages_path = api_packages_path,
+        loading_value = if_indexing_source_code("static", loading_value),
+        visibility = visibility,
+    )
+
+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", "")
+
+    #    dir_path = "tensorflow/%s" % dir_path
+    return dir_path.replace("/", ".").strip(".")
+
+def _api_gen_rule_impl(ctx):
+    api_gen_binary_target = ctx.attr.api_gen_binary_target[DefaultInfo].files_to_run.executable
+    flags = [ctx.expand_location(flag) for flag in ctx.attr.flags]
+    variables = {"@D": ctx.genfiles_dir.path + "/" + ctx.label.package}
+    flags = [ctx.expand_make_variables("tf_api_version", flag, variables) for flag in flags]
+    loading = ctx.expand_make_variables("TF_API_INIT_LOADING", ctx.attr.loading_value, {})
+    output_paths = [f.path for f in ctx.outputs.outs]
+
+    # Generate file containing the list of outputs
+    # Without this, the command will be too long (even when executed in a shell script)
+    params = ctx.actions.declare_file(ctx.attr.name + ".params")
+    ctx.actions.write(params, ";".join(output_paths))
+
+    # Convert output_paths to the list of corresponding modules for the further testing
+    if ctx.attr.api_packages_path:
+        output_modules = sorted([
+            _get_module_by_path(path, ctx.genfiles_dir.path)
+            for path in output_paths
+            if "__init__.py" in path
+        ])
+        api_packages = ctx.actions.declare_file(ctx.attr.api_packages_path.name)
+        ctx.actions.write(api_packages, "\n".join(output_modules))
+
+    cmd = _make_cmd(api_gen_binary_target, flags, loading, [params.path])
+    ctx.actions.run_shell(
+        inputs = ctx.files.srcs + [params],
+        outputs = ctx.outputs.outs,
+        mnemonic = "TensorFlowApiGen",
+        tools = [api_gen_binary_target],
+        use_default_shell_env = True,
+        command = cmd,
+    )
+
+# Note: if only one output_paths is provided, api_gen_binary_target assumes it is a file to be read
+def _make_cmd(api_gen_binary_target, flags, loading, output_paths):
+    binary = api_gen_binary_target.path
+    flags = flags + ["--loading=" + loading]
+    return " ".join([binary] + flags + output_paths)
+
+# To prevent compiling the C++ code twice, we only want to build `api_gen_binary_target`
+# for the target platform and not the execution platform.
+# To achieve this without causing confusion with source dependencies (e.g. putting api_gen_binary_target in srcs of the genrule),
+# we use a custom rule to execute the command line for generating the API files.
+# See https://github.com/tensorflow/tensorflow/issues/60167
+# To not break internal cross-platform builds, we only set `cfg` to `target` for the OSS build.
+api_gen_rule = rule(
+    implementation = _api_gen_rule_impl,
+    attrs = {
+        "outs": attr.output_list(mandatory = True),
+        "srcs": attr.label_list(allow_files = True),
+        "flags": attr.string_list(),
+        "api_gen_binary_target": attr.label(executable = True, cfg = if_oss("target", "exec"), mandatory = True),
+        "loading_value": attr.string(mandatory = True),
+        "api_packages_path": attr.output(),
+    },
+)

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

@@ -0,0 +1,183 @@
+"""TensorFlow V2 API __init__.py files."""
+
+# keep sorted
+TENSORFLOW_API_INIT_FILES = [
+    # BEGIN GENERATED FILES
+    "__init__.py",
+    "__internal__/__init__.py",
+    "__internal__/autograph/__init__.py",
+    "__internal__/decorator/__init__.py",
+    "__internal__/dispatch/__init__.py",
+    "__internal__/distribute/__init__.py",
+    "__internal__/distribute/combinations/__init__.py",
+    "__internal__/distribute/interim/__init__.py",
+    "__internal__/distribute/multi_process_runner/__init__.py",
+    "__internal__/eager_context/__init__.py",
+    "__internal__/feature_column/__init__.py",
+    "__internal__/function/__init__.py",
+    "__internal__/graph_util/__init__.py",
+    "__internal__/mixed_precision/__init__.py",
+    "__internal__/monitoring/__init__.py",
+    "__internal__/nest/__init__.py",
+    "__internal__/ops/__init__.py",
+    "__internal__/smart_cond/__init__.py",
+    "__internal__/test/__init__.py",
+    "__internal__/test/combinations/__init__.py",
+    "__internal__/tf2/__init__.py",
+    "__internal__/train/__init__.py",
+    "__internal__/types/__init__.py",
+    "__internal__/types/data/__init__.py",
+    "__internal__/saved_model/__init__.py",
+    "__internal__/saved_model/load/__init__.py",
+    "__internal__/tracking/__init__.py",
+    "__operators__/__init__.py",
+    "audio/__init__.py",
+    "autograph/__init__.py",
+    "autograph/experimental/__init__.py",
+    "autodiff/__init__.py",
+    "bitwise/__init__.py",
+    "compat/__init__.py",
+    "config/__init__.py",
+    "config/experimental/__init__.py",
+    "config/optimizer/__init__.py",
+    "config/threading/__init__.py",
+    "data/__init__.py",
+    "data/experimental/__init__.py",
+    "data/experimental/service/__init__.py",
+    "debugging/__init__.py",
+    "debugging/experimental/__init__.py",
+    "distribute/__init__.py",
+    "distribute/cluster_resolver/__init__.py",
+    "distribute/coordinator/__init__.py",
+    "distribute/experimental/__init__.py",
+    "distribute/experimental/coordinator/__init__.py",
+    "distribute/experimental/partitioners/__init__.py",
+    "distribute/experimental/rpc/__init__.py",
+    "dtypes/__init__.py",
+    "dtypes/experimental/__init__.py",
+    "errors/__init__.py",
+    "experimental/__init__.py",
+    "experimental/extension_type/__init__.py",
+    "experimental/dtensor/__init__.py",
+    "experimental/numpy/__init__.py",
+    "experimental/numpy/random/__init__.py",
+    "experimental/tensorrt/__init__.py",
+    "experimental/dlpack/__init__.py",
+    "feature_column/__init__.py",
+    "io/gfile/__init__.py",
+    "graph_util/__init__.py",
+    "image/__init__.py",
+    "io/__init__.py",
+    "queue/__init__.py",
+    "linalg/__init__.py",
+    "linalg/experimental/__init__.py",
+    "lite/__init__.py",
+    "lite/experimental/__init__.py",
+    "lite/experimental/authoring/__init__.py",
+    "lookup/__init__.py",
+    "lookup/experimental/__init__.py",
+    "math/__init__.py",
+    "math/special/__init__.py",
+    "mlir/__init__.py",
+    "mlir/experimental/__init__.py",
+    "nest/__init__.py",
+    "nn/__init__.py",
+    "nn/experimental/__init__.py",
+    "profiler/__init__.py",
+    "profiler/experimental/__init__.py",
+    "profiler/experimental/client/__init__.py",
+    "profiler/experimental/server/__init__.py",
+    "quantization/__init__.py",
+    "quantization/experimental/__init__.py",
+    "ragged/__init__.py",
+    "random/__init__.py",
+    "random/experimental/__init__.py",
+    "raw_ops/__init__.py",
+    "saved_model/__init__.py",
+    "saved_model/experimental/__init__.py",
+    "sets/__init__.py",
+    "signal/__init__.py",
+    "sparse/__init__.py",
+    "strings/__init__.py",
+    "summary/__init__.py",
+    "summary/experimental/__init__.py",
+    "sysconfig/__init__.py",
+    "test/__init__.py",
+    "test/experimental/__init__.py",
+    "tpu/experimental/embedding/__init__.py",
+    "tpu/experimental/__init__.py",
+    "tpu/__init__.py",
+    "train/__init__.py",
+    "train/experimental/__init__.py",
+    "types/__init__.py",
+    "types/experimental/__init__.py",
+    "types/experimental/distributed/__init__.py",
+    "version/__init__.py",
+    "xla/__init__.py",
+    "xla/experimental/__init__.py",
+    # END GENERATED FILES
+]
+
+KERAS_API_INIT_FILES = [
+    "__init__.py",
+    "keras/__init__.py",
+    "keras/__internal__/__init__.py",
+    "keras/__internal__/backend/__init__.py",
+    "keras/__internal__/models/__init__.py",
+    "keras/__internal__/losses/__init__.py",
+    "keras/__internal__/layers/__init__.py",
+    "keras/__internal__/optimizers/__init__.py",
+    "keras/__internal__/utils/__init__.py",
+    "keras/activations/__init__.py",
+    "keras/applications/__init__.py",
+    "keras/applications/densenet/__init__.py",
+    "keras/applications/efficientnet/__init__.py",
+    "keras/applications/imagenet_utils/__init__.py",
+    "keras/applications/inception_resnet_v2/__init__.py",
+    "keras/applications/inception_v3/__init__.py",
+    "keras/applications/mobilenet/__init__.py",
+    "keras/applications/mobilenet_v2/__init__.py",
+    "keras/applications/mobilenet_v3/__init__.py",
+    "keras/applications/nasnet/__init__.py",
+    "keras/applications/resnet/__init__.py",
+    "keras/applications/resnet_v2/__init__.py",
+    "keras/applications/resnet50/__init__.py",
+    "keras/applications/vgg16/__init__.py",
+    "keras/applications/vgg19/__init__.py",
+    "keras/applications/xception/__init__.py",
+    "keras/backend/__init__.py",
+    "keras/callbacks/__init__.py",
+    "keras/callbacks/experimental/__init__.py",
+    "keras/constraints/__init__.py",
+    "keras/datasets/__init__.py",
+    "keras/datasets/boston_housing/__init__.py",
+    "keras/datasets/cifar10/__init__.py",
+    "keras/datasets/cifar100/__init__.py",
+    "keras/datasets/fashion_mnist/__init__.py",
+    "keras/datasets/imdb/__init__.py",
+    "keras/datasets/mnist/__init__.py",
+    "keras/datasets/reuters/__init__.py",
+    "keras/estimator/__init__.py",
+    "keras/experimental/__init__.py",
+    # Placeholder for internal API
+    "keras/initializers/__init__.py",
+    "keras/layers/__init__.py",
+    "keras/layers/experimental/__init__.py",
+    "keras/layers/experimental/preprocessing/__init__.py",
+    "keras/losses/__init__.py",
+    "keras/metrics/__init__.py",
+    "keras/mixed_precision/__init__.py",
+    "keras/premade/__init__.py",
+    "keras/models/__init__.py",
+    "keras/optimizers/__init__.py",
+    "keras/optimizers/experimental/__init__.py",
+    "keras/optimizers/schedules/__init__.py",
+    "keras/preprocessing/__init__.py",
+    "keras/preprocessing/image/__init__.py",
+    "keras/preprocessing/sequence/__init__.py",
+    "keras/preprocessing/text/__init__.py",
+    "keras/regularizers/__init__.py",
+    "keras/utils/__init__.py",
+    "keras/utils/experimental/__init__.py",
+    "keras/utils/legacy/__init__.py",
+]

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

@@ -0,0 +1,166 @@
+"""TensorFlow V1 API __init__.py files."""
+
+# keep sorted
+TENSORFLOW_API_INIT_FILES_V1 = [
+    # BEGIN GENERATED FILES
+    "__init__.py",
+    "__internal__/__init__.py",
+    "__internal__/types/__init__.py",
+    "__internal__/types/data/__init__.py",
+    "app/__init__.py",
+    "audio/__init__.py",
+    "autograph/__init__.py",
+    "autograph/experimental/__init__.py",
+    "bitwise/__init__.py",
+    "compat/__init__.py",
+    "config/__init__.py",
+    "config/experimental/__init__.py",
+    "config/optimizer/__init__.py",
+    "config/threading/__init__.py",
+    "data/__init__.py",
+    "data/experimental/__init__.py",
+    "data/experimental/service/__init__.py",
+    "debugging/__init__.py",
+    "debugging/experimental/__init__.py",
+    "distribute/__init__.py",
+    "distribute/cluster_resolver/__init__.py",
+    "distribute/experimental/__init__.py",
+    "distributions/__init__.py",
+    "dtypes/__init__.py",
+    "dtypes/experimental/__init__.py",
+    "errors/__init__.py",
+    "experimental/__init__.py",
+    "experimental/extension_type/__init__.py",
+    "feature_column/__init__.py",
+    "gfile/__init__.py",
+    "io/gfile/__init__.py",
+    "graph_util/__init__.py",
+    "image/__init__.py",
+    "io/__init__.py",
+    "queue/__init__.py",
+    "initializers/__init__.py",
+    "layers/__init__.py",
+    "linalg/__init__.py",
+    "linalg/experimental/__init__.py",
+    "lite/__init__.py",
+    "lite/constants/__init__.py",
+    "lite/experimental/__init__.py",
+    "lite/experimental/authoring/__init__.py",
+    "logging/__init__.py",
+    "lookup/__init__.py",
+    "lookup/experimental/__init__.py",
+    "losses/__init__.py",
+    "manip/__init__.py",
+    "math/__init__.py",
+    "math/special/__init__.py",
+    "metrics/__init__.py",
+    "mixed_precision/__init__.py",
+    "mixed_precision/experimental/__init__.py",
+    "mlir/__init__.py",
+    "mlir/experimental/__init__.py",
+    "nest/__init__.py",
+    "nn/__init__.py",
+    "nn/experimental/__init__.py",
+    "nn/rnn_cell/__init__.py",
+    "profiler/__init__.py",
+    "python_io/__init__.py",
+    "quantization/__init__.py",
+    "quantization/experimental/__init__.py",
+    "ragged/__init__.py",
+    "random/__init__.py",
+    "random/experimental/__init__.py",
+    "raw_ops/__init__.py",
+    "resource_loader/__init__.py",
+    "strings/__init__.py",
+    "saved_model/__init__.py",
+    "saved_model/builder/__init__.py",
+    "saved_model/constants/__init__.py",
+    "saved_model/experimental/__init__.py",
+    "saved_model/loader/__init__.py",
+    "saved_model/main_op/__init__.py",
+    "saved_model/signature_constants/__init__.py",
+    "saved_model/signature_def_utils/__init__.py",
+    "saved_model/tag_constants/__init__.py",
+    "saved_model/utils/__init__.py",
+    "sets/__init__.py",
+    "signal/__init__.py",
+    "sparse/__init__.py",
+    "spectral/__init__.py",
+    "summary/__init__.py",
+    "sysconfig/__init__.py",
+    "test/__init__.py",
+    "test/experimental/__init__.py",
+    "tpu/experimental/embedding/__init__.py",
+    "tpu/experimental/__init__.py",
+    "tpu/__init__.py",
+    "train/__init__.py",
+    "train/experimental/__init__.py",
+    "train/queue_runner/__init__.py",
+    "types/__init__.py",
+    "types/experimental/__init__.py",
+    "user_ops/__init__.py",
+    "version/__init__.py",
+    "xla/__init__.py",
+    "xla/experimental/__init__.py",
+    # END GENERATED FILES
+]
+
+KERAS_API_INIT_FILES_V1 = [
+    "__init__.py",
+    "keras/__init__.py",
+    "keras/__internal__/__init__.py",
+    "keras/__internal__/layers/__init__.py",
+    "keras/__internal__/legacy/__init__.py",
+    "keras/__internal__/legacy/layers/__init__.py",
+    "keras/__internal__/legacy/layers/experimental/__init__.py",
+    "keras/__internal__/legacy/rnn_cell/__init__.py",
+    "keras/activations/__init__.py",
+    "keras/applications/__init__.py",
+    "keras/applications/densenet/__init__.py",
+    "keras/applications/efficientnet/__init__.py",
+    "keras/applications/imagenet_utils/__init__.py",
+    "keras/applications/inception_resnet_v2/__init__.py",
+    "keras/applications/inception_v3/__init__.py",
+    "keras/applications/mobilenet/__init__.py",
+    "keras/applications/mobilenet_v2/__init__.py",
+    "keras/applications/mobilenet_v3/__init__.py",
+    "keras/applications/nasnet/__init__.py",
+    "keras/applications/resnet/__init__.py",
+    "keras/applications/resnet_v2/__init__.py",
+    "keras/applications/resnet50/__init__.py",
+    "keras/applications/vgg16/__init__.py",
+    "keras/applications/vgg19/__init__.py",
+    "keras/applications/xception/__init__.py",
+    "keras/backend/__init__.py",
+    "keras/callbacks/__init__.py",
+    "keras/callbacks/experimental/__init__.py",
+    "keras/constraints/__init__.py",
+    "keras/datasets/__init__.py",
+    "keras/datasets/boston_housing/__init__.py",
+    "keras/datasets/cifar10/__init__.py",
+    "keras/datasets/cifar100/__init__.py",
+    "keras/datasets/fashion_mnist/__init__.py",
+    "keras/datasets/imdb/__init__.py",
+    "keras/datasets/mnist/__init__.py",
+    "keras/datasets/reuters/__init__.py",
+    "keras/estimator/__init__.py",
+    "keras/experimental/__init__.py",
+    "keras/initializers/__init__.py",
+    "keras/layers/__init__.py",
+    "keras/layers/experimental/__init__.py",
+    "keras/layers/experimental/preprocessing/__init__.py",
+    "keras/losses/__init__.py",
+    "keras/metrics/__init__.py",
+    "keras/mixed_precision/__init__.py",
+    "keras/models/__init__.py",
+    "keras/optimizers/__init__.py",
+    "keras/optimizers/schedules/__init__.py",
+    "keras/premade/__init__.py",
+    "keras/preprocessing/__init__.py",
+    "keras/preprocessing/image/__init__.py",
+    "keras/preprocessing/sequence/__init__.py",
+    "keras/preprocessing/text/__init__.py",
+    "keras/regularizers/__init__.py",
+    "keras/utils/__init__.py",
+    "keras/utils/legacy/__init__.py",
+]

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

@@ -0,0 +1,854 @@
+# Copyright 2017 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.
+# =============================================================================
+"""Generates and prints out imports and constants for new TensorFlow python api."""
+import argparse
+import collections
+import importlib
+import os
+import sys
+
+from tensorflow.python.tools.api.generator import doc_srcs
+from tensorflow.python.util import tf_decorator
+from tensorflow.python.util import tf_export
+
+API_ATTRS = tf_export.API_ATTRS
+API_ATTRS_V1 = tf_export.API_ATTRS_V1
+
+_LAZY_LOADING = False
+_API_VERSIONS = [1, 2]
+_COMPAT_MODULE_TEMPLATE = 'compat.v%d'
+_SUBCOMPAT_MODULE_TEMPLATE = 'compat.v%d.compat.v%d'
+_COMPAT_MODULE_PREFIX = 'compat.v'
+_DEFAULT_PACKAGE = 'tensorflow.python'
+_GENFILES_DIR_SUFFIX = 'genfiles/'
+_SYMBOLS_TO_SKIP_EXPLICITLY = {
+    # Overrides __getattr__, so that unwrapping tf_decorator
+    # would have side effects.
+    'tensorflow.python.platform.flags.FLAGS'
+}
+_GENERATED_FILE_HEADER = """# This file is MACHINE GENERATED! Do not edit.
+# Generated by: tensorflow/python/tools/api/generator/create_python_api.py script.
+\"\"\"%s
+\"\"\"
+
+import sys as _sys
+
+"""
+_GENERATED_FILE_FOOTER = ''
+_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)
+"""
+_LAZY_LOADING_MODULE_TEXT_TEMPLATE = """
+# Inform pytype that this module is dynamically populated (b/111239204).
+_HAS_DYNAMIC_ATTRIBUTES = True
+_PUBLIC_APIS = {
+%s
+}
+"""
+
+
+class SymbolExposedTwiceError(Exception):
+  """Raised when different symbols are exported with the same name."""
+  pass
+
+
+def get_canonical_import(import_set):
+  """Obtain one single import from a set of possible sources of a symbol.
+
+  One symbol might come from multiple places as it is being imported and
+  reexported. To simplify API changes, we always use the same import for the
+  same module, and give preference based on higher priority and alphabetical
+  ordering.
+
+  Args:
+    import_set: (set) Imports providing the same symbol. This is a set of tuples
+      in the form (import, priority). We want to pick an import with highest
+      priority.
+
+  Returns:
+    A module name to import
+  """
+  # We use the fact that list sorting is stable, so first we convert the set to
+  # a sorted list of the names and then we resort this list to move elements
+  # not in core tensorflow to the end.
+  # Here we sort by priority (higher preferred) and then  alphabetically by
+  # import string.
+  import_list = sorted(
+      import_set,
+      key=lambda imp_and_priority: (-imp_and_priority[1], imp_and_priority[0]))
+  return import_list[0][0]
+
+
+class _ModuleInitCodeBuilder(object):
+  """Builds a map from module name to imports included in that module."""
+
+  def __init__(self,
+               output_package,
+               api_version,
+               lazy_loading=_LAZY_LOADING,
+               use_relative_imports=False):
+    self._output_package = output_package
+    # Maps API module to API symbol name to set of tuples of the form
+    # (module name, priority).
+    # The same symbol can be imported from multiple locations. Higher
+    # "priority" indicates that import location is preferred over others.
+    self._module_imports = collections.defaultdict(
+        lambda: collections.defaultdict(set))
+    self._dest_import_to_id = collections.defaultdict(int)
+    # Names that start with underscore in the root module.
+    self._underscore_names_in_root = set()
+    self._api_version = api_version
+    # Controls whether or not exported symbols are lazily loaded or statically
+    # imported.
+    self._lazy_loading = lazy_loading
+    self._use_relative_imports = use_relative_imports
+
+  def _check_already_imported(self, symbol_id, api_name):
+    if (api_name in self._dest_import_to_id and
+        symbol_id != self._dest_import_to_id[api_name] and symbol_id != -1):
+      raise SymbolExposedTwiceError(
+          f'Trying to export multiple symbols with same name: {api_name}')
+    self._dest_import_to_id[api_name] = symbol_id
+
+  def add_import(self, symbol, source_module_name, source_name,
+                 dest_module_name, dest_name):
+    """Adds this import to module_imports.
+
+    Args:
+      symbol: TensorFlow Python symbol.
+      source_module_name: (string) Module to import from.
+      source_name: (string) Name of the symbol to import.
+      dest_module_name: (string) Module name to add import to.
+      dest_name: (string) Import the symbol using this name.
+
+    Raises:
+      SymbolExposedTwiceError: Raised when an import with the same
+        dest_name has already been added to dest_module_name.
+    """
+    # modules_with_exports.py is only used during API generation and
+    # won't be available when actually importing tensorflow.
+    if source_module_name.endswith('python.modules_with_exports'):
+      source_module_name = symbol.__module__
+    import_str = self.format_import(source_module_name, source_name, dest_name)
+
+    # Check if we are trying to expose two different symbols with same name.
+    full_api_name = dest_name
+    if dest_module_name:
+      full_api_name = dest_module_name + '.' + full_api_name
+    symbol_id = -1 if not symbol else id(symbol)
+    self._check_already_imported(symbol_id, full_api_name)
+
+    if not dest_module_name and dest_name.startswith('_'):
+      self._underscore_names_in_root.add(dest_name)
+
+    # The same symbol can be available in multiple modules.
+    # We store all possible ways of importing this symbol and later pick just
+    # one.
+    priority = 0
+    if symbol:
+      # Give higher priority to source module if it matches
+      # symbol's original module.
+      if hasattr(symbol, '__module__'):
+        priority = int(source_module_name == symbol.__module__)
+      # Give higher priority if symbol name matches its __name__.
+      if hasattr(symbol, '__name__'):
+        priority += int(source_name == symbol.__name__)
+    self._module_imports[dest_module_name][full_api_name].add(
+        (import_str, priority))
+
+  def _import_submodules(self):
+    """Add imports for all destination modules in self._module_imports."""
+    # Import all required modules in their parent modules.
+    # For e.g. if we import 'foo.bar.Value'. Then, we also
+    # import 'bar' in 'foo'.
+    imported_modules = set(self._module_imports.keys())
+    for module in imported_modules:
+      if not module:
+        continue
+      module_split = module.split('.')
+      parent_module = ''  # we import submodules in their parent_module
+
+      for submodule_index in range(len(module_split)):
+        if submodule_index > 0:
+          submodule = module_split[submodule_index - 1]
+          parent_module += '.' + submodule if parent_module else submodule
+        import_from = self._output_package
+        if self._lazy_loading:
+          import_from += '.' + '.'.join(module_split[:submodule_index + 1])
+          self.add_import(
+              symbol=None,
+              source_module_name='',
+              source_name=import_from,
+              dest_module_name=parent_module,
+              dest_name=module_split[submodule_index])
+        else:
+          if self._use_relative_imports:
+            import_from = '.'
+          elif submodule_index > 0:
+            import_from += '.' + '.'.join(module_split[:submodule_index])
+          self.add_import(
+              symbol=None,
+              source_module_name=import_from,
+              source_name=module_split[submodule_index],
+              dest_module_name=parent_module,
+              dest_name=module_split[submodule_index])
+
+  def build(self):
+    """Get a map from destination module to __init__.py code for that module.
+
+    Returns:
+      A dictionary where
+        key: (string) destination module (for e.g. tf or tf.consts).
+        value: (string) text that should be in __init__.py files for
+          corresponding modules.
+    """
+    self._import_submodules()
+    module_text_map = {}
+    footer_text_map = {}
+    for dest_module, dest_name_to_imports in self._module_imports.items():
+      # Sort all possible imports for a symbol and pick the first one.
+      imports_list = [
+          get_canonical_import(imports)
+          for _, imports in dest_name_to_imports.items()
+      ]
+      if self._lazy_loading:
+        module_text_map[
+            dest_module] = _LAZY_LOADING_MODULE_TEXT_TEMPLATE % '\n'.join(
+                sorted(imports_list))
+      else:
+        module_text_map[dest_module] = '\n'.join(sorted(imports_list))
+
+    # Expose exported symbols with underscores in root module since we import
+    # from it using * import. Don't need this for lazy_loading because the
+    # underscore symbols are already included in __all__ when passed in and
+    # handled by TFModuleWrapper.
+    root_module_footer = ''
+    if not self._lazy_loading:
+      underscore_names_str = ', '.join(
+          '\'%s\'' % name for name in sorted(self._underscore_names_in_root))
+
+      root_module_footer = """
+_names_with_underscore = [%s]
+__all__ = [_s for _s in dir() if not _s.startswith('_')]
+__all__.extend([_s for _s in _names_with_underscore])
+""" % underscore_names_str
+
+    # Add module wrapper if we need to print deprecation messages
+    # or if we use lazy loading.
+    if self._api_version == 1 or self._lazy_loading:
+      for dest_module, _ in self._module_imports.items():
+        deprecation = 'False'
+        has_lite = 'False'
+        if self._api_version == 1:  # Add 1.* deprecations.
+          if not dest_module.startswith(_COMPAT_MODULE_PREFIX):
+            deprecation = 'True'
+        # Workaround to make sure not load lite from lite/__init__.py
+        if (not dest_module and 'lite' in self._module_imports and
+            self._lazy_loading):
+          has_lite = 'True'
+        if self._lazy_loading:
+          public_apis_name = '_PUBLIC_APIS'
+        else:
+          public_apis_name = 'None'
+        footer_text_map[dest_module] = _DEPRECATION_FOOTER % (
+            dest_module, public_apis_name, deprecation, has_lite)
+
+    return module_text_map, footer_text_map, root_module_footer
+
+  def format_import(self, source_module_name, source_name, dest_name):
+    """Formats import statement.
+
+    Args:
+      source_module_name: (string) Source module to import from.
+      source_name: (string) Source symbol name to import.
+      dest_name: (string) Destination alias name.
+
+    Returns:
+      An import statement string.
+    """
+    if self._lazy_loading:
+      return "  '%s': ('%s', '%s')," % (dest_name, source_module_name,
+                                        source_name)
+    else:
+      if source_module_name:
+        if source_name == dest_name:
+          return 'from %s import %s' % (source_module_name, source_name)
+        else:
+          return 'from %s import %s as %s' % (source_module_name, source_name,
+                                              dest_name)
+      else:
+        if source_name == dest_name:
+          return 'import %s' % source_name
+        else:
+          return 'import %s as %s' % (source_name, dest_name)
+
+  def get_destination_modules(self):
+    return set(self._module_imports.keys())
+
+  def copy_imports(self, from_dest_module, to_dest_module):
+    self._module_imports[to_dest_module] = (
+        self._module_imports[from_dest_module].copy())
+
+
+def add_nested_compat_imports(module_builder, compat_api_versions,
+                              output_package):
+  """Adds compat.vN.compat.vK modules to module builder.
+
+  To avoid circular imports, we want to add __init__.py files under
+  compat.vN.compat.vK and under compat.vN.compat.vK.compat. For all other
+  imports, we point to corresponding modules under compat.vK.
+
+  Args:
+    module_builder: `_ModuleInitCodeBuilder` instance.
+    compat_api_versions: Supported compatibility versions.
+    output_package: Base output python package where generated API will be
+      added.
+  """
+  imported_modules = module_builder.get_destination_modules()
+
+  # Copy over all imports in compat.vK to compat.vN.compat.vK and
+  # all imports in compat.vK.compat to compat.vN.compat.vK.compat.
+  for v in compat_api_versions:
+    for sv in compat_api_versions:
+      subcompat_module = _SUBCOMPAT_MODULE_TEMPLATE % (v, sv)
+      compat_module = _COMPAT_MODULE_TEMPLATE % sv
+      module_builder.copy_imports(compat_module, subcompat_module)
+      module_builder.copy_imports('%s.compat' % compat_module,
+                                  '%s.compat' % subcompat_module)
+
+  # Prefixes of modules under compatibility packages, for e.g. "compat.v1.".
+  compat_prefixes = tuple(
+      _COMPAT_MODULE_TEMPLATE % v + '.' for v in compat_api_versions)
+
+  # Above, we only copied function, class and constant imports. Here
+  # we also add imports for child modules.
+  for imported_module in imported_modules:
+    if not imported_module.startswith(compat_prefixes):
+      continue
+    module_split = imported_module.split('.')
+
+    # Handle compat.vN.compat.vK.compat.foo case. That is,
+    # import compat.vK.compat.foo in compat.vN.compat.vK.compat.
+    if len(module_split) > 3 and module_split[2] == 'compat':
+      src_module = '.'.join(module_split[:3])
+      src_name = module_split[3]
+      assert src_name != 'v1' and src_name != 'v2', imported_module
+    else:  # Handle compat.vN.compat.vK.foo case.
+      src_module = '.'.join(module_split[:2])
+      src_name = module_split[2]
+      if src_name == 'compat':
+        continue  # compat.vN.compat.vK.compat is handled separately
+
+    for compat_api_version in compat_api_versions:
+      module_builder.add_import(
+          symbol=None,
+          source_module_name='%s.%s' % (output_package, src_module),
+          source_name=src_name,
+          dest_module_name='compat.v%d.%s' % (compat_api_version, src_module),
+          dest_name=src_name)
+
+
+def _get_name_and_module(full_name):
+  """Split full_name into module and short name.
+
+  Args:
+    full_name: Full name of symbol that includes module.
+
+  Returns:
+    Full module name and short symbol name.
+  """
+  name_segments = full_name.split('.')
+  return '.'.join(name_segments[:-1]), name_segments[-1]
+
+
+def _join_modules(module1, module2):
+  """Concatenate 2 module components.
+
+  Args:
+    module1: First module to join.
+    module2: Second module to join.
+
+  Returns:
+    Given two modules aaa.bbb and ccc.ddd, returns a joined
+    module aaa.bbb.ccc.ddd.
+  """
+  if not module1:
+    return module2
+  if not module2:
+    return module1
+  return '%s.%s' % (module1, module2)
+
+
+def add_imports_for_symbol(module_code_builder,
+                           symbol,
+                           source_module_name,
+                           source_name,
+                           api_name,
+                           api_version,
+                           output_module_prefix=''):
+  """Add imports for the given symbol to `module_code_builder`.
+
+  Args:
+    module_code_builder: `_ModuleInitCodeBuilder` instance.
+    symbol: A symbol.
+    source_module_name: Module that we can import the symbol from.
+    source_name: Name we can import the symbol with.
+    api_name: API name. Currently, must be `tensorflow`.
+    api_version: API version.
+    output_module_prefix: Prefix to prepend to destination module.
+  """
+  if api_version == 1:
+    names_attr = API_ATTRS_V1[api_name].names
+    constants_attr = API_ATTRS_V1[api_name].constants
+  else:
+    names_attr = API_ATTRS[api_name].names
+    constants_attr = API_ATTRS[api_name].constants
+
+  # If symbol is _tf_api_constants attribute, then add the constants.
+  if source_name == constants_attr:
+    for exports, name in symbol:
+      for export in exports:
+        dest_module, dest_name = _get_name_and_module(export)
+        dest_module = _join_modules(output_module_prefix, dest_module)
+        module_code_builder.add_import(None, source_module_name, name,
+                                       dest_module, dest_name)
+
+  # If symbol has _tf_api_names attribute, then add import for it.
+  if (hasattr(symbol, '__dict__') and names_attr in symbol.__dict__):
+
+    # Generate import statements for symbols.
+    for export in getattr(symbol, names_attr):  # pylint: disable=protected-access
+      dest_module, dest_name = _get_name_and_module(export)
+      dest_module = _join_modules(output_module_prefix, dest_module)
+      module_code_builder.add_import(symbol, source_module_name, source_name,
+                                     dest_module, dest_name)
+
+
+def get_api_init_text(packages,
+                      packages_to_ignore,
+                      output_package,
+                      api_name,
+                      api_version,
+                      compat_api_versions=None,
+                      lazy_loading=_LAZY_LOADING,
+                      use_relative_imports=False):
+  """Get a map from destination module to __init__.py code for that module.
+
+  Args:
+    packages: Base python packages containing python with target tf_export
+      decorators.
+    packages_to_ignore: python packages to be ignored when checking for
+      tf_export decorators.
+    output_package: Base output python package where generated API will be
+      added.
+    api_name: API you want to generate Currently, only `tensorflow`.
+    api_version: API version you want to generate (1 or 2).
+    compat_api_versions: Additional API versions to generate under compat/
+      directory.
+    lazy_loading: Boolean flag. If True, a lazy loading `__init__.py` file is
+      produced and if `False`, static imports are used.
+    use_relative_imports: True if we should use relative imports when importing
+      submodules.
+
+  Returns:
+    A dictionary where
+      key: (string) destination module (for e.g. tf or tf.consts).
+      value: (string) text that should be in __init__.py files for
+        corresponding modules.
+  """
+  if compat_api_versions is None:
+    compat_api_versions = []
+  module_code_builder = _ModuleInitCodeBuilder(output_package, api_version,
+                                               lazy_loading,
+                                               use_relative_imports)
+
+  # Traverse over everything imported above. Specifically,
+  # we want to traverse over TensorFlow Python modules.
+
+  def in_packages(m):
+    return any(package in m for package in packages)
+
+  for module in list(sys.modules.values()):
+    # Only look at tensorflow modules.
+    if (not module or not hasattr(module, '__name__') or
+        module.__name__ is None or not in_packages(module.__name__)):
+      continue
+    if packages_to_ignore and any([p for p in packages_to_ignore
+                                   if p in module.__name__]):
+      continue
+
+    # Do not generate __init__.py files for contrib modules for now.
+    if (('.contrib.' in module.__name__ or module.__name__.endswith('.contrib'))
+        and '.lite' not in module.__name__):
+      continue
+
+    for module_contents_name in dir(module):
+      if (module.__name__ + '.' +
+          module_contents_name in _SYMBOLS_TO_SKIP_EXPLICITLY):
+        continue
+      attr = getattr(module, module_contents_name)
+      _, attr = tf_decorator.unwrap(attr)
+
+      add_imports_for_symbol(module_code_builder, attr, module.__name__,
+                             module_contents_name, api_name, api_version)
+      for compat_api_version in compat_api_versions:
+        add_imports_for_symbol(module_code_builder, attr, module.__name__,
+                               module_contents_name, api_name,
+                               compat_api_version,
+                               _COMPAT_MODULE_TEMPLATE % compat_api_version)
+
+  if compat_api_versions:
+    add_nested_compat_imports(module_code_builder, compat_api_versions,
+                              output_package)
+  return module_code_builder.build()
+
+
+def get_module(dir_path, relative_to_dir):
+  """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 get_module_docstring(module_name, package, api_name):
+  """Get docstring for the given module.
+
+  This method looks for docstring in the following order:
+  1. Checks if module has a docstring specified in doc_srcs.
+  2. Checks if module has a docstring source module specified
+     in doc_srcs. If it does, gets docstring from that module.
+  3. Checks if module with module_name exists under base package.
+     If it does, gets docstring from that module.
+  4. Returns a default docstring.
+
+  Args:
+    module_name: module name relative to tensorflow (excluding 'tensorflow.'
+      prefix) to get a docstring for.
+    package: Base python package containing python with target tf_export
+      decorators.
+    api_name: API you want to generate Currently, only `tensorflow`.
+
+  Returns:
+    One-line docstring to describe the module.
+  """
+  # Get the same module doc strings for any version. That is, for module
+  # 'compat.v1.foo' we can get docstring from module 'foo'.
+  for version in _API_VERSIONS:
+    compat_prefix = _COMPAT_MODULE_TEMPLATE % version
+    if module_name.startswith(compat_prefix):
+      module_name = module_name[len(compat_prefix):].strip('.')
+
+  # Module under base package to get a docstring from.
+  docstring_module_name = module_name
+
+  doc_sources = doc_srcs.get_doc_sources(api_name)
+
+  if module_name in doc_sources:
+    docsrc = doc_sources[module_name]
+    if docsrc.docstring:
+      return docsrc.docstring
+    if docsrc.docstring_module_name:
+      docstring_module_name = docsrc.docstring_module_name
+
+  if package != 'tf_keras':
+    docstring_module_name = package + '.' + docstring_module_name
+  if (docstring_module_name in sys.modules and
+      sys.modules[docstring_module_name].__doc__):
+    return sys.modules[docstring_module_name].__doc__
+
+  return 'Public API for tf.%s namespace.' % module_name
+
+
+def create_primary_api_files(output_files,
+                             packages,
+                             packages_to_ignore,
+                             root_init_template,
+                             output_dir,
+                             output_package,
+                             api_name,
+                             api_version,
+                             compat_api_versions,
+                             compat_init_templates,
+                             lazy_loading=_LAZY_LOADING,
+                             use_relative_imports=False):
+  """Creates __init__.py files for the Python API.
+
+  Args:
+    output_files: List of __init__.py file paths to create.
+    packages: Base python packages containing python with target tf_export
+      decorators.
+    packages_to_ignore: python packages to be ignored when checking for
+      tf_export decorators.
+    root_init_template: Template for top-level __init__.py file. "# API IMPORTS
+      PLACEHOLDER" comment in the template file will be replaced with imports.
+    output_dir: output API root directory.
+    output_package: Base output package where generated API will be added.
+    api_name: API you want to generate Currently, only `tensorflow`.
+    api_version: API version to generate (`v1` or `v2`).
+    compat_api_versions: Additional API versions to generate in compat/
+      subdirectory.
+    compat_init_templates: List of templates for top level compat init files in
+      the same order as compat_api_versions.
+    lazy_loading: Boolean flag. If True, a lazy loading `__init__.py` file is
+      produced and if `False`, static imports are used.
+    use_relative_imports: True if we should use relative imports when import
+      submodules.
+
+  Raises:
+    ValueError: if output_files list is missing a required file.
+  """
+  module_name_to_file_path = {}
+  for output_file in output_files:
+    module_name = get_module(os.path.dirname(output_file), output_dir)
+    module_name_to_file_path[module_name] = os.path.normpath(output_file)
+
+  # Create file for each expected output in genrule.
+  for module, file_path in module_name_to_file_path.items():
+    if not os.path.isdir(os.path.dirname(file_path)):
+      os.makedirs(os.path.dirname(file_path))
+    open(file_path, 'a').close()
+
+  (
+      module_text_map,
+      deprecation_footer_map,
+      root_module_footer,
+  ) = get_api_init_text(packages, packages_to_ignore, output_package, api_name,
+                        api_version, compat_api_versions, lazy_loading,
+                        use_relative_imports)
+
+  # Add imports to output files.
+  missing_output_files = []
+  # Root modules are "" and "compat.v*".
+  root_module = ''
+  compat_module_to_template = {
+      _COMPAT_MODULE_TEMPLATE % v: t
+      for v, t in zip(compat_api_versions, compat_init_templates)
+  }
+  for v in compat_api_versions:
+    compat_module_to_template.update({
+        _SUBCOMPAT_MODULE_TEMPLATE % (v, vs): t
+        for vs, t in zip(compat_api_versions, compat_init_templates)
+    })
+
+  for module, text in module_text_map.items():
+    # Make sure genrule output file list is in sync with API exports.
+    if module not in module_name_to_file_path:
+      module_file_path = '"%s/__init__.py"' % (module.replace('.', '/'))
+      missing_output_files.append(module_file_path)
+      continue
+
+    contents = ''
+    if module == root_module and root_init_template:
+      # Read base init file for root module
+      with open(root_init_template, 'r') as root_init_template_file:
+        contents = root_init_template_file.read()
+        contents = contents.replace('# API IMPORTS PLACEHOLDER', text)
+        contents = contents.replace('# __all__ PLACEHOLDER', root_module_footer)
+    elif module in compat_module_to_template:
+      # Read base init file for compat module
+      with open(compat_module_to_template[module], 'r') as init_template_file:
+        contents = init_template_file.read()
+        contents = contents.replace('# API IMPORTS PLACEHOLDER', text)
+    else:
+      contents = (
+          _GENERATED_FILE_HEADER %
+          get_module_docstring(module, packages[0], api_name) + text +
+          _GENERATED_FILE_FOOTER)
+    if module in deprecation_footer_map:
+      if '# WRAPPER_PLACEHOLDER' in contents:
+        contents = contents.replace('# WRAPPER_PLACEHOLDER',
+                                    deprecation_footer_map[module])
+      else:
+        contents += deprecation_footer_map[module]
+    with open(module_name_to_file_path[module], 'w') as fp:
+      fp.write(contents)
+
+  if missing_output_files:
+    missing_files = ',\n'.join(sorted(missing_output_files))
+    raise ValueError(
+        f'Missing outputs for genrule:\n{missing_files}. Be sure to add these '
+        'targets to tensorflow/python/tools/api/generator/api_init_files_v1.bzl'
+        ' and tensorflow/python/tools/api/generator/api_init_files.bzl '
+        '(tensorflow repo), or tf_keras/api/api_init_files.bzl (tf_keras repo)')
+
+
+def create_proxy_api_files(output_files,
+                           proxy_module_root,
+                           output_dir):
+  """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_path in output_files:
+    module = get_module(os.path.dirname(file_path), output_dir)
+    if not os.path.isdir(os.path.dirname(file_path)):
+      os.makedirs(os.path.dirname(file_path))
+    contents = f'from {proxy_module_root}.{module} import *'
+    with open(file_path, 'w') as fp:
+      fp.write(contents)
+
+
+def main():
+  parser = argparse.ArgumentParser()
+  parser.add_argument(
+      'outputs',
+      metavar='O',
+      type=str,
+      nargs='+',
+      help='If a single file is passed in, then we assume it contains a '
+      'semicolon-separated list of Python files that we expect this script to '
+      'output. If multiple files are passed in, then we assume output files '
+      'are listed directly as arguments.')
+  parser.add_argument(
+      '--packages',
+      default=_DEFAULT_PACKAGE,
+      type=str,
+      help='Base packages that import modules containing the target tf_export '
+      'decorators.')
+  parser.add_argument(
+      '--packages_to_ignore',
+      default='',
+      type=str,
+      help='Packages to exclude from the api generation. This is used to hide '
+      'certain packages from this script when multiple copy of code exists, '
+      'eg tf_keras. It is useful to avoid the SymbolExposedTwiceError.'
+      )
+  parser.add_argument(
+      '--root_init_template',
+      default='',
+      type=str,
+      help='Template for top level __init__.py file. '
+      '"#API IMPORTS PLACEHOLDER" comment will be replaced with imports.')
+  parser.add_argument(
+      '--apidir',
+      type=str,
+      required=True,
+      help='Directory where generated output files are placed. '
+      'gendir should be a prefix of apidir. Also, apidir '
+      'should be a prefix of every directory in outputs.')
+  parser.add_argument(
+      '--apiname',
+      required=True,
+      type=str,
+      choices=API_ATTRS.keys(),
+      help='The API you want to generate.')
+  parser.add_argument(
+      '--apiversion',
+      default=2,
+      type=int,
+      choices=_API_VERSIONS,
+      help='The API version you want to generate.')
+  parser.add_argument(
+      '--compat_apiversions',
+      default=[],
+      type=int,
+      action='append',
+      help='Additional versions to generate in compat/ subdirectory. '
+      'If set to 0, then no additional version would be generated.')
+  parser.add_argument(
+      '--compat_init_templates',
+      default=[],
+      type=str,
+      action='append',
+      help='Templates for top-level __init__ files under compat modules. '
+      'The list of init file templates must be in the same order as '
+      'list of versions passed with compat_apiversions.')
+  parser.add_argument(
+      '--output_package',
+      default='tensorflow',
+      type=str,
+      help='Root output package.')
+  parser.add_argument(
+      '--loading',
+      default='default',
+      type=str,
+      choices=['lazy', 'static', 'default'],
+      help='Controls how the generated __init__.py file loads the exported '
+      'symbols. \'lazy\' means the symbols are loaded when first used. '
+      '\'static\' means all exported symbols are loaded in the '
+      '__init__.py file. \'default\' uses the value of the '
+      '_LAZY_LOADING constant in create_python_api.py.')
+  parser.add_argument(
+      '--use_relative_imports',
+      default=False,
+      type=bool,
+      help='Whether to import submodules using relative imports or absolute '
+      'imports')
+  parser.add_argument(
+      '--proxy_module_root',
+      default=None,
+      type=str,
+      help='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.')
+  args = parser.parse_args()
+
+  if len(args.outputs) == 1:
+    # If we only get a single argument, then it must be a file containing
+    # list of outputs.
+    with open(args.outputs[0]) as output_list_file:
+      outputs = [line.strip() for line in output_list_file.read().split(';')]
+  else:
+    outputs = args.outputs
+
+  # Populate `sys.modules` with modules containing tf_export().
+  packages = args.packages.split(',')
+  for package in packages:
+    importlib.import_module(package)
+  packages_to_ignore = args.packages_to_ignore.split(',')
+
+  # Determine if the modules shall be loaded lazily or statically.
+  if args.loading == 'default':
+    lazy_loading = _LAZY_LOADING
+  elif args.loading == 'lazy':
+    lazy_loading = True
+  elif args.loading == 'static':
+    lazy_loading = False
+  else:
+    # This should never happen (tm).
+    raise ValueError(f'Invalid value for --loading flag: {args.loading}. Must '
+                     'be one of lazy, static, default.')
+  if args.proxy_module_root is None:
+    create_primary_api_files(outputs, packages, packages_to_ignore,
+                             args.root_init_template, args.apidir,
+                             args.output_package, args.apiname, args.apiversion,
+                             args.compat_apiversions,
+                             args.compat_init_templates, lazy_loading,
+                             args.use_relative_imports)
+  else:
+    create_proxy_api_files(outputs, args.proxy_module_root, args.apidir)
+
+
+if __name__ == '__main__':
+  main()

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

@@ -0,0 +1,181 @@
+# Copyright 2017 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.
+# =============================================================================
+"""Tests for create_python_api."""
+
+import os
+import sys
+import types
+
+from tensorflow.python.platform import test
+from tensorflow.python.tools.api.generator import create_python_api
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export('test_op', 'test_op1', 'test.test_op2')
+def test_op():
+  pass
+
+
+@tf_export('test1.foo', v1=['test.foo'])
+def deprecated_test_op():
+  pass
+
+
+@tf_export('TestClass', 'NewTestClass')
+class TestClass(object):
+  pass
+
+
+_TEST_CONSTANT = 5
+_MODULE_NAME = 'tensorflow.python.test_module'
+
+
+class CreatePythonApiTest(test.TestCase):
+
+  def setUp(self):
+    # Add fake op to a module that has 'tensorflow' in the name.
+    sys.modules[_MODULE_NAME] = types.ModuleType(_MODULE_NAME)
+    setattr(sys.modules[_MODULE_NAME], 'test_op', test_op)
+    setattr(sys.modules[_MODULE_NAME], 'deprecated_test_op', deprecated_test_op)
+    setattr(sys.modules[_MODULE_NAME], 'TestClass', TestClass)
+    test_op.__module__ = _MODULE_NAME
+    TestClass.__module__ = _MODULE_NAME
+    tf_export('consts._TEST_CONSTANT').export_constant(
+        _MODULE_NAME, '_TEST_CONSTANT')
+
+  def tearDown(self):
+    del sys.modules[_MODULE_NAME]
+
+  def testFunctionImportIsAdded(self):
+    imports, _, _ = create_python_api.get_api_init_text(
+        packages=[create_python_api._DEFAULT_PACKAGE],
+        packages_to_ignore=[],
+        output_package='tensorflow',
+        api_name='tensorflow',
+        api_version=1)
+    if create_python_api._LAZY_LOADING:
+      expected_import = (
+          '\'test_op1\': '
+          '(\'tensorflow.python.test_module\','
+          ' \'test_op\')')
+    else:
+      expected_import = (
+          'from tensorflow.python.test_module '
+          'import test_op as test_op1')
+    self.assertTrue(
+        expected_import in str(imports),
+        msg='%s not in %s' % (expected_import, str(imports)))
+
+    if create_python_api._LAZY_LOADING:
+      expected_import = (
+          '\'test_op\': '
+          '(\'tensorflow.python.test_module\','
+          ' \'test_op\')')
+    else:
+      expected_import = (
+          'from tensorflow.python.test_module '
+          'import test_op')
+    self.assertTrue(
+        expected_import in str(imports),
+        msg='%s not in %s' % (expected_import, str(imports)))
+    # Also check that compat.v1 is not added to imports.
+    self.assertFalse('compat.v1' in imports,
+                     msg='compat.v1 in %s' % str(imports.keys()))
+
+  def testClassImportIsAdded(self):
+    imports, _, _ = create_python_api.get_api_init_text(
+        packages=[create_python_api._DEFAULT_PACKAGE],
+        packages_to_ignore=[],
+        output_package='tensorflow',
+        api_name='tensorflow',
+        api_version=2)
+    if create_python_api._LAZY_LOADING:
+      expected_import = (
+          '\'NewTestClass\':'
+          ' (\'tensorflow.python.test_module\','
+          ' \'TestClass\')')
+    else:
+      expected_import = (
+          'from tensorflow.python.test_module '
+          'import TestClass')
+    self.assertTrue(
+        'TestClass' in str(imports),
+        msg='%s not in %s' % (expected_import, str(imports)))
+
+  def testConstantIsAdded(self):
+    imports, _, _ = create_python_api.get_api_init_text(
+        packages=[create_python_api._DEFAULT_PACKAGE],
+        packages_to_ignore=[],
+        output_package='tensorflow',
+        api_name='tensorflow',
+        api_version=1)
+    if create_python_api._LAZY_LOADING:
+      expected = ('\'_TEST_CONSTANT\':'
+                  ' (\'tensorflow.python.test_module\','
+                  ' \'_TEST_CONSTANT\')')
+    else:
+      expected = ('from tensorflow.python.test_module '
+                  'import _TEST_CONSTANT')
+    self.assertTrue(expected in str(imports),
+                    msg='%s not in %s' % (expected, str(imports)))
+
+  def testCompatModuleIsAdded(self):
+    imports, _, _ = create_python_api.get_api_init_text(
+        packages=[create_python_api._DEFAULT_PACKAGE],
+        packages_to_ignore=[],
+        output_package='tensorflow',
+        api_name='tensorflow',
+        api_version=2,
+        compat_api_versions=[1])
+    self.assertTrue('compat.v1' in imports,
+                    msg='compat.v1 not in %s' % str(imports.keys()))
+    self.assertTrue('compat.v1.test' in imports,
+                    msg='compat.v1.test not in %s' % str(imports.keys()))
+
+  def testNestedCompatModulesAreAdded(self):
+    imports, _, _ = create_python_api.get_api_init_text(
+        packages=[create_python_api._DEFAULT_PACKAGE],
+        packages_to_ignore=[],
+        output_package='tensorflow',
+        api_name='tensorflow',
+        api_version=2,
+        compat_api_versions=[1, 2])
+    self.assertIn('compat.v1.compat.v1', imports,
+                  msg='compat.v1.compat.v1 not in %s' % str(imports.keys()))
+    self.assertIn('compat.v1.compat.v2', imports,
+                  msg='compat.v1.compat.v2 not in %s' % str(imports.keys()))
+    self.assertIn('compat.v2.compat.v1', imports,
+                  msg='compat.v2.compat.v1 not in %s' % str(imports.keys()))
+    self.assertIn('compat.v2.compat.v2', imports,
+                  msg='compat.v2.compat.v2 not in %s' % str(imports.keys()))
+
+  def testProxyAPIFileIsGenerated(self):
+    save_dir = self.get_temp_dir()
+    proxy_module_root = 'tf_keras.api._v2'
+    module = 'keras.losses'
+    module_dir = module.replace('.', '/')
+    proxy_file = os.path.join(save_dir, module_dir, '__init__.py')
+    expected_imports = [f'from {proxy_module_root}.{module} import *']
+    create_python_api.create_proxy_api_files([proxy_file],
+                                             proxy_module_root,
+                                             save_dir)
+    self.assertTrue(os.path.exists(proxy_file))
+    with open(proxy_file, 'r') as f:
+      lines = f.readlines()
+    self.assertCountEqual(expected_imports, lines)
+
+
+if __name__ == '__main__':
+  test.main()

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

@@ -0,0 +1,144 @@
+# Copyright 2018 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.
+# ==============================================================================
+"""Specifies sources of doc strings for API modules."""
+from tensorflow.python.util import tf_export
+
+
+class DocSource(object):
+  """Specifies docstring source for a module.
+
+  Only one of docstring or docstring_module_name should be set.
+  * If docstring is set, then we will use this docstring when
+    for the module.
+  * If docstring_module_name is set, then we will copy the docstring
+    from docstring source module.
+  """
+
+  def __init__(self, docstring=None, docstring_module_name=None):
+    self.docstring = docstring
+    self.docstring_module_name = docstring_module_name
+
+    if self.docstring is not None and self.docstring_module_name is not None:
+      raise ValueError('Only one of `docstring` or `docstring_module_name` can '
+                       'be set.')
+
+
+_TENSORFLOW_DOC_SOURCES = {
+    'app':
+        DocSource(docstring='Import router for absl.app.'),
+    'bitwise':
+        DocSource(docstring_module_name='ops.bitwise_ops'),
+    'compat':
+        DocSource(docstring_module_name='util.compat'),
+    'distribute':
+        DocSource(docstring_module_name='distribute'),
+    'distributions': DocSource(
+        docstring='Core module for TensorFlow distribution objects and helpers.'
+    ),
+    'errors':
+        DocSource(docstring_module_name='framework.errors'),
+    'experimental.numpy':
+        DocSource(docstring_module_name='ops.numpy_ops'),
+    'gfile':
+        DocSource(docstring='Import router for file_io.'),
+    'graph_util':
+        DocSource(docstring_module_name='framework.graph_util'),
+    'image':
+        DocSource(docstring_module_name='ops.image_ops'),
+    'linalg':
+        DocSource(docstring_module_name='ops.linalg_ops'),
+    'logging':
+        DocSource(docstring_module_name='ops.logging_ops'),
+    'losses': DocSource(
+        docstring=(
+            'Loss operations for use in neural networks. Note: All the losses'
+            ' are added to the `GraphKeys.LOSSES` collection by default.'
+        )
+    ),
+    'manip':
+        DocSource(docstring_module_name='ops.manip_ops'),
+    'math':
+        DocSource(docstring_module_name='ops.math_ops'),
+    'metrics':
+        DocSource(docstring='Evaluation-related metrics.'),
+    'nest':
+        DocSource(docstring_module_name='util.nest'),
+    'nn':
+        DocSource(docstring_module_name='ops.nn_ops'),
+    'nn.rnn_cell':
+        DocSource(docstring='Module for constructing RNN Cells.'),
+    'python_io':
+        DocSource(docstring_module_name='lib.io.python_io'),
+    'ragged':
+        DocSource(docstring_module_name='ops.ragged'),
+    'resource_loader':
+        DocSource(docstring='Resource management library.'),
+    'sets':
+        DocSource(docstring_module_name='ops.sets'),
+    'signal':
+        DocSource(docstring_module_name='ops.signal'),
+    'sparse':
+        DocSource(docstring_module_name='ops.sparse_ops'),
+    'strings':
+        DocSource(docstring_module_name='ops.string_ops'),
+    'summary': DocSource(
+        docstring=(
+            'Operations for writing summary data, for use in analysis and'
+            ' visualization. See the [Summaries and'
+            ' TensorBoard](https://www.tensorflow.org/guide/summaries_and_tensorboard)'
+            ' guide.'
+        )
+    ),
+    'sysconfig':
+        DocSource(docstring='System configuration library.'),
+    'test':
+        DocSource(docstring='Testing.'),
+    'train': DocSource(
+        docstring=(
+            'Support for training models. See the'
+            ' [Training](https://tensorflow.org/api_guides/python/train) guide.'
+        )
+    ),
+}
+
+_ESTIMATOR_DOC_SOURCES = {
+    'estimator': DocSource(
+        docstring_module_name='estimator_lib'),
+    'estimator.export': DocSource(
+        docstring_module_name='export.export_lib'),
+    'estimator.inputs': DocSource(
+        docstring_module_name='inputs.inputs'),
+}
+
+_KERAS_DOC_SOURCES = {
+    'keras.optimizers.experimental':
+        DocSource(docstring_module_name='optimizers.optimizer_experimental')
+}
+
+
+def get_doc_sources(api_name):
+  """Get a map from module to a DocSource object.
+
+  Args:
+    api_name: API you want to generate (e.g. `tensorflow` or `estimator`).
+
+  Returns:
+    Map from module name to DocSource object.
+  """
+  if api_name == tf_export.TENSORFLOW_API_NAME:
+    return _TENSORFLOW_DOC_SOURCES
+  if api_name == tf_export.KERAS_API_NAME:
+    return _KERAS_DOC_SOURCES
+  return {}

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

@@ -0,0 +1,79 @@
+# Copyright 2018 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.
+# =============================================================================
+"""Tests for tensorflow.python.tools.api.generator.doc_srcs."""
+
+import argparse
+import importlib
+import sys
+
+from tensorflow.python.platform import test
+from tensorflow.python.tools.api.generator import doc_srcs
+
+
+FLAGS = None
+
+
+class DocSrcsTest(test.TestCase):
+
+  def testModulesAreValidAPIModules(self):
+    for module_name in doc_srcs.get_doc_sources(FLAGS.api_name):
+      # Convert module_name to corresponding __init__.py file path.
+      file_path = module_name.replace('.', '/')
+      if file_path:
+        file_path += '/'
+      file_path += '__init__.py'
+
+      self.assertIn(
+          file_path, FLAGS.outputs,
+          msg='%s is not a valid API module' % module_name)
+
+  def testHaveDocstringOrDocstringModule(self):
+    for module_name, docsrc in doc_srcs.get_doc_sources(FLAGS.api_name).items():
+      self.assertFalse(
+          docsrc.docstring and docsrc.docstring_module_name,
+          msg=('%s contains DocSource has both a docstring and a '
+               'docstring_module_name. Only one of "docstring" or '
+               '"docstring_module_name" should be set.') % (module_name))
+
+  def testDocstringModulesAreValidModules(self):
+    for _, docsrc in doc_srcs.get_doc_sources(FLAGS.api_name).items():
+      if docsrc.docstring_module_name:
+        doc_module_name = '.'.join([
+            FLAGS.package, docsrc.docstring_module_name])
+        self.assertIn(
+            doc_module_name, sys.modules,
+            msg=('docsources_module %s is not a valid module under %s.' %
+                 (docsrc.docstring_module_name, FLAGS.package)))
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+  parser.add_argument(
+      'outputs', metavar='O', type=str, nargs='+',
+      help='create_python_api output files.')
+  parser.add_argument(
+      '--package', type=str,
+      help='Base package that imports modules containing the target tf_export '
+           'decorators.')
+  parser.add_argument(
+      '--api_name', type=str,
+      help='API name: tensorflow')
+  FLAGS, unparsed = parser.parse_known_args()
+
+  importlib.import_module(FLAGS.package)
+
+  # Now update argv, so that unittest library does not get confused.
+  sys.argv = [sys.argv[0]] + unparsed
+  test.main()

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

@@ -0,0 +1,195 @@
+# Copyright 2018 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.
+# =============================================================================
+"""Tests for api_init_files.bzl and api_init_files_v1.bzl."""
+import argparse
+import importlib
+import sys
+
+from tensorflow.python.platform import resource_loader
+from tensorflow.python.platform import test
+from tensorflow.python.util import tf_decorator
+
+
+def _traverse_packages(packages):
+  for package in packages:
+    importlib.import_module(package)
+
+
+def _get_module_from_symbol(symbol):
+  if '.' not in symbol:
+    return ''
+  return '.'.join(symbol.split('.')[:-1])
+
+
+def _get_modules(package, attr_name, constants_attr_name):
+  """Get list of TF API modules.
+
+  Args:
+    package: We only look at modules that contain package in the name.
+    attr_name: Attribute set on TF symbols that contains API names.
+    constants_attr_name: Attribute set on TF modules that contains
+      API constant names.
+
+  Returns:
+    Set of TensorFlow API modules.
+  """
+  modules = set()
+  # TODO(annarev): split up the logic in create_python_api.py so that
+  #   it can be reused in this test.
+  for module in list(sys.modules.values()):
+    if (not module or not hasattr(module, '__name__') or
+        package not in module.__name__):
+      continue
+
+    for module_contents_name in dir(module):
+      attr = getattr(module, module_contents_name)
+      _, attr = tf_decorator.unwrap(attr)
+
+      # Add modules to _tf_api_constants attribute.
+      if module_contents_name == constants_attr_name:
+        for exports, _ in attr:
+          modules.update(
+              [_get_module_from_symbol(export) for export in exports])
+        continue
+
+      # Add modules for _tf_api_names attribute.
+      if (hasattr(attr, '__dict__') and attr_name in attr.__dict__):
+        modules.update([
+            _get_module_from_symbol(export)
+            for export in getattr(attr, attr_name)])
+  return modules
+
+
+def _get_files_set(path, start_tag, end_tag):
+  """Get set of file paths from the given file.
+
+  Args:
+    path: Path to file. File at `path` is expected to contain a list of paths
+      where entire list starts with `start_tag` and ends with `end_tag`. List
+      must be comma-separated and each path entry must be surrounded by double
+      quotes.
+    start_tag: String that indicates start of path list.
+    end_tag: String that indicates end of path list.
+
+  Returns:
+    List of string paths.
+  """
+  with open(path, 'r') as f:
+    contents = f.read()
+    start = contents.find(start_tag) + len(start_tag) + 1
+    end = contents.find(end_tag)
+    contents = contents[start:end]
+    file_paths = [
+        file_path.strip().strip('"') for file_path in contents.split(',')]
+    return set(file_path for file_path in file_paths if file_path)
+
+
+def _module_to_paths(module):
+  """Get all API __init__.py file paths for the given module.
+
+  Args:
+    module: Module to get file paths for.
+
+  Returns:
+    List of paths for the given module. For e.g. module foo.bar
+    requires 'foo/__init__.py' and 'foo/bar/__init__.py'.
+  """
+  submodules = []
+  module_segments = module.split('.')
+  for i in range(len(module_segments)):
+    submodules.append('.'.join(module_segments[:i+1]))
+  paths = []
+  for submodule in submodules:
+    if not submodule:
+      paths.append('__init__.py')
+      continue
+    paths.append('%s/__init__.py' % (submodule.replace('.', '/')))
+  return paths
+
+
+class OutputInitFilesTest(test.TestCase):
+  """Test that verifies files that list paths for TensorFlow API."""
+
+  def _validate_paths_for_modules(
+      self, actual_paths, expected_paths, file_to_update_on_error):
+    """Validates that actual_paths match expected_paths.
+
+    Args:
+      actual_paths: */__init__.py file paths listed in file_to_update_on_error.
+      expected_paths: */__init__.py file paths that we need to create for
+        TensorFlow API.
+      file_to_update_on_error: File that contains list of */__init__.py files.
+        We include it in error message printed if the file list needs to be
+        updated.
+    """
+    self.assertTrue(actual_paths)
+    self.assertTrue(expected_paths)
+    missing_paths = expected_paths - actual_paths
+    extra_paths = actual_paths - expected_paths
+
+    # Surround paths with quotes so that they can be copy-pasted
+    # from error messages as strings.
+    missing_paths = ['\'%s\'' % path for path in missing_paths]
+    extra_paths = ['\'%s\'' % path for path in extra_paths]
+
+    self.assertFalse(
+        missing_paths,
+        'Please add %s to %s.' % (
+            ',\n'.join(sorted(missing_paths)), file_to_update_on_error))
+    self.assertFalse(
+        extra_paths,
+        'Redundant paths, please remove %s in %s.' % (
+            ',\n'.join(sorted(extra_paths)), file_to_update_on_error))
+
+  def test_V2_init_files(self):
+    modules = _get_modules(
+        'tensorflow', '_tf_api_names', '_tf_api_constants')
+    file_path = resource_loader.get_path_to_datafile(
+        'api_init_files.bzl')
+    paths = _get_files_set(
+        file_path, '# BEGIN GENERATED FILES', '# END GENERATED FILES')
+    module_paths = set(
+        f for module in modules for f in _module_to_paths(module))
+    self._validate_paths_for_modules(
+        paths, module_paths, file_to_update_on_error=file_path)
+
+  def test_V1_init_files(self):
+    modules = _get_modules(
+        'tensorflow', '_tf_api_names_v1', '_tf_api_constants_v1')
+    file_path = resource_loader.get_path_to_datafile(
+        'api_init_files_v1.bzl')
+    paths = _get_files_set(
+        file_path, '# BEGIN GENERATED FILES', '# END GENERATED FILES')
+    module_paths = set(
+        f for module in modules for f in _module_to_paths(module))
+    self._validate_paths_for_modules(
+        paths, module_paths, file_to_update_on_error=file_path)
+
+
+if __name__ == '__main__':
+  parser = argparse.ArgumentParser()
+  parser.add_argument(
+      '--packages',
+      type=str,
+      default='',
+      help='Comma separated list of packages to traverse.')
+  FLAGS, unparsed = parser.parse_known_args()
+
+  # Traverse packages that define APIs.
+  _traverse_packages(FLAGS.packages.split(','))
+
+  # Now update argv, so that unittest library does not get confused.
+  sys.argv = [sys.argv[0]] + unparsed
+  test.main()

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()