wehub-resource-sync
362 lines
12 KiB
Python
362 lines
12 KiB
Python
# 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.
|
|
# ==============================================================================
|
|
"""A tool to generate api_docs for TensorFlow2.
|
|
|
|
```
|
|
python generate2.py --output_dir=/tmp/out
|
|
```
|
|
|
|
Requires a local installation of `tensorflow_docs`:
|
|
|
|
```
|
|
pip install git+https://github.com/tensorflow/docs
|
|
```
|
|
"""
|
|
import contextlib
|
|
import pathlib
|
|
import textwrap
|
|
from typing import NamedTuple
|
|
|
|
from absl import app
|
|
from absl import flags
|
|
from packaging import version
|
|
import tensorflow as tf
|
|
from tensorflow_docs.api_generator import doc_controls
|
|
from tensorflow_docs.api_generator import doc_generator_visitor
|
|
from tensorflow_docs.api_generator import generate_lib
|
|
from tensorflow_docs.api_generator.pretty_docs import base_page
|
|
from tensorflow_docs.api_generator.pretty_docs import module_page
|
|
import yaml
|
|
|
|
from tensorflow.python.framework import ops
|
|
from tensorflow.python.util import tf_export
|
|
from tensorflow.python.util import tf_inspect
|
|
|
|
if version.parse(tf.__version__) >= version.parse("2.14-dev"):
|
|
from tensorflow.python.util.pywrap_xla_ops import get_gpu_kernel_names # pylint: disable=g-import-not-at-top
|
|
|
|
# Caution: the google and oss versions of this import are different.
|
|
import base_dir # pylint: disable=g-import-not-at-top
|
|
|
|
# pylint: disable=g-import-not-at-top
|
|
try:
|
|
from tensorflow.python.types import doc_typealias
|
|
_EXTRA_DOCS = getattr(doc_typealias, "_EXTRA_DOCS", {})
|
|
del doc_typealias
|
|
except ImportError:
|
|
_EXTRA_DOCS = {}
|
|
# pylint: enable=g-import-not-at-top
|
|
|
|
# `tf` has an `__all__` that doesn't list important things like `keras`.
|
|
# The doc generator recognizes `__all__` as the list of public symbols.
|
|
# So patch `tf.__all__` to list everything.
|
|
tf.__all__ = [item_name for item_name, value in tf_inspect.getmembers(tf)]
|
|
|
|
# tf_export generated two copies of the module objects.
|
|
# This will just list compat.v2 as an alias for tf. Close enough, let's not
|
|
# duplicate all the module skeleton files.
|
|
tf.compat.v2 = tf
|
|
|
|
tf.losses = tf.keras.losses
|
|
tf.metrics = tf.keras.metrics
|
|
tf.optimizers = tf.keras.optimizers
|
|
tf.initializers = tf.keras.initializers
|
|
|
|
MIN_NUM_FILES_EXPECTED = 2000
|
|
FLAGS = flags.FLAGS
|
|
|
|
flags.DEFINE_string(
|
|
"code_url_prefix",
|
|
"/code/stable/tensorflow",
|
|
"A url to prepend to code paths when creating links to defining code")
|
|
|
|
flags.DEFINE_string("output_dir", "/tmp/out",
|
|
"A directory, where the docs will be output to.")
|
|
|
|
flags.DEFINE_bool("search_hints", True,
|
|
"Include meta-data search hints at the top of each file.")
|
|
|
|
flags.DEFINE_string(
|
|
"site_path", "",
|
|
"The path prefix (up to `.../api_docs/python`) used in the "
|
|
"`_toc.yaml` and `_redirects.yaml` files")
|
|
|
|
_PRIVATE_MAP = {
|
|
"tf": ["python", "core", "compiler", "examples", "tools", "contrib"],
|
|
# There's some aliasing between the compats and v1/2s, so it's easier to
|
|
# block by name and location than by deleting, or hiding objects.
|
|
"tf.compat.v1.compat": ["v1", "v2"],
|
|
"tf.compat.v2.compat": ["v1", "v2"]
|
|
}
|
|
|
|
tf.__doc__ = """
|
|
## TensorFlow
|
|
|
|
```
|
|
pip install tensorflow
|
|
```
|
|
"""
|
|
|
|
try:
|
|
tf.estimator.Estimator = doc_controls.inheritable_header(textwrap.dedent("""\
|
|
Warning: TensorFlow 2.15 included the final release of the `tf-estimator`
|
|
package. Estimators will not be available in TensorFlow 2.16 or after. See the
|
|
[migration guide](https://www.tensorflow.org/guide/migrate/migrating_estimator)
|
|
for more information about how to convert off of Estimators."
|
|
"""))(tf.estimator.Estimator)
|
|
except AttributeError:
|
|
pass
|
|
|
|
|
|
class RawOpsPageInfo(module_page.ModulePageInfo):
|
|
"""Generates a custom page for `tf.raw_ops`."""
|
|
|
|
DEFAULT_BUILDER_CLASS = base_page.TemplatePageBuilder
|
|
|
|
def build(self):
|
|
# Skip the ModulePage implementation, which doesn't use a template.
|
|
content = base_page.PageInfo.build(self)
|
|
|
|
if version.parse(tf.__version__) >= version.parse("2.14-dev"):
|
|
raw_ops_doc = self.generate_raw_ops_doc_ge_214()
|
|
else:
|
|
raw_ops_doc = self.generate_raw_ops_doc_lt_214()
|
|
|
|
return "\n".join([content, raw_ops_doc])
|
|
|
|
def generate_raw_ops_doc_lt_214(self):
|
|
"""Generates docs for `tf.raw_ops`."""
|
|
del self
|
|
|
|
warning = textwrap.dedent("""\n
|
|
Note: `tf.raw_ops` provides direct/low level access to all TensorFlow ops.
|
|
See [the RFC](https://github.com/tensorflow/community/blob/master/rfcs/20181225-tf-raw-ops.md)
|
|
for details. Unless you are library writer, you likely do not need to use
|
|
these ops directly.""")
|
|
|
|
table_header = textwrap.dedent("""
|
|
|
|
| Op Name | Has Gradient |
|
|
|---------|:------------:|""")
|
|
|
|
parts = [warning, table_header]
|
|
|
|
for op_name in sorted(dir(tf.raw_ops)):
|
|
try:
|
|
ops._gradient_registry.lookup(op_name) # pylint: disable=protected-access
|
|
has_gradient = "\N{HEAVY CHECK MARK}\N{VARIATION SELECTOR-16}"
|
|
except LookupError:
|
|
has_gradient = "\N{CROSS MARK}"
|
|
|
|
if not op_name.startswith("_"):
|
|
path = pathlib.Path("/") / FLAGS.site_path / "tf/raw_ops" / op_name
|
|
path = path.with_suffix(".md")
|
|
link = ('<a id={op_name} href="{path}">{op_name}</a>').format(
|
|
op_name=op_name, path=str(path))
|
|
parts.append("| {link} | {has_gradient} |".format(
|
|
link=link, has_gradient=has_gradient))
|
|
|
|
return "\n".join(parts)
|
|
|
|
def generate_raw_ops_doc_ge_214(self):
|
|
"""Generates docs for `tf.raw_ops`."""
|
|
del self
|
|
|
|
warning = textwrap.dedent("""\n
|
|
Note: `tf.raw_ops` provides direct/low level access to all TensorFlow ops.
|
|
See [the RFC](https://github.com/tensorflow/community/blob/master/rfcs/20181225-tf-raw-ops.md)
|
|
for details. Unless you are library writer, you likely do not need to use
|
|
these ops directly.""")
|
|
|
|
table_header = textwrap.dedent("""
|
|
|
|
| Op Name | Has Gradient | GPU XLA Support |
|
|
|---------|:------------:|:---------------:|""")
|
|
|
|
parts = [warning, table_header]
|
|
xla_compiled_ops = get_gpu_kernel_names()
|
|
for op_name in sorted(dir(tf.raw_ops)):
|
|
try:
|
|
ops._gradient_registry.lookup(op_name) # pylint: disable=protected-access
|
|
has_gradient = "\N{HEAVY CHECK MARK}\N{VARIATION SELECTOR-16}"
|
|
except LookupError:
|
|
has_gradient = "\N{CROSS MARK}"
|
|
is_xla_compilable = "\N{CROSS MARK}"
|
|
if op_name in xla_compiled_ops:
|
|
is_xla_compilable = "\N{HEAVY CHECK MARK}\N{VARIATION SELECTOR-16}"
|
|
|
|
if not op_name.startswith("_"):
|
|
path = pathlib.Path("/") / FLAGS.site_path / "tf/raw_ops" / op_name
|
|
path = path.with_suffix(".md")
|
|
link = ('<a id={op_name} href="{path}">{op_name}</a>').format(
|
|
op_name=op_name, path=str(path)
|
|
)
|
|
parts.append(
|
|
"| {link} | {has_gradient} | {is_xla_compilable} |".format(
|
|
link=link,
|
|
has_gradient=has_gradient,
|
|
is_xla_compilable=is_xla_compilable,
|
|
)
|
|
)
|
|
|
|
return "\n".join(parts)
|
|
|
|
|
|
# The doc generator isn't aware of tf_export.
|
|
# So prefix the score tuples with -1 when this is the canonical name, +1
|
|
# otherwise. The generator chooses the name with the lowest score.
|
|
class TfExportAwareVisitor(doc_generator_visitor.DocGeneratorVisitor):
|
|
"""A `tf_export`, `keras_export` and `estimator_export` aware doc_visitor."""
|
|
|
|
class TfNameScore(NamedTuple):
|
|
canonical_score: int
|
|
name_score: doc_generator_visitor.DocGeneratorVisitor.NameScore
|
|
|
|
def _score_name(self, path: doc_generator_visitor.ApiPath) -> TfNameScore:
|
|
name = ".".join(path)
|
|
all_exports = [
|
|
tf_export.TENSORFLOW_API_NAME,
|
|
tf_export.KERAS_API_NAME,
|
|
]
|
|
|
|
try:
|
|
all_exports.append(tf_export.ESTIMATOR_API_NAME)
|
|
except AttributeError:
|
|
pass
|
|
|
|
canonical = None
|
|
for api_name in all_exports:
|
|
try:
|
|
canonical = tf_export.get_canonical_name_for_symbol(
|
|
self._index[name], api_name=api_name)
|
|
except AttributeError:
|
|
canonical = None
|
|
if canonical is not None:
|
|
break
|
|
|
|
canonical_score = 1
|
|
if canonical is not None and name == "tf." + canonical:
|
|
canonical_score = -1
|
|
|
|
return self.TfNameScore(canonical_score, super()._score_name(path))
|
|
|
|
|
|
def build_docs(output_dir, code_url_prefix, search_hints):
|
|
"""Build api docs for tensorflow v2.
|
|
|
|
Args:
|
|
output_dir: A string path, where to put the files.
|
|
code_url_prefix: prefix for "Defined in" links.
|
|
search_hints: Bool. Include meta-data search hints at the top of each file.
|
|
"""
|
|
output_dir = pathlib.Path(output_dir)
|
|
site_path = pathlib.Path("/", FLAGS.site_path)
|
|
|
|
doc_controls.set_deprecated(tf.compat.v1)
|
|
try:
|
|
doc_controls.set_deprecated(tf.estimator)
|
|
except AttributeError:
|
|
pass
|
|
doc_controls.set_deprecated(tf.feature_column)
|
|
doc_controls.set_deprecated(tf.keras.preprocessing)
|
|
|
|
# The custom page will be used for raw_ops.md not the one generated above.
|
|
doc_controls.set_custom_page_builder_cls(tf.raw_ops, RawOpsPageInfo)
|
|
|
|
# Hide raw_ops from search.
|
|
for name, obj in tf_inspect.getmembers(tf.raw_ops):
|
|
if not name.startswith("_"):
|
|
doc_controls.hide_from_search(obj)
|
|
|
|
for cls in [tf.Module, tf.keras.layers.Layer, tf.keras.optimizers.Optimizer]:
|
|
doc_controls.decorate_all_class_attributes(
|
|
decorator=doc_controls.do_not_doc_in_subclasses,
|
|
cls=cls,
|
|
skip=["__init__"])
|
|
|
|
do_not_document = ["tf.__internal__",
|
|
"tf.keras.__internal__",
|
|
"tf.keras.wrappers",
|
|
"tf.__operators__",
|
|
"tf.tools",
|
|
"tf.compat.v1.pywrap_tensorflow",
|
|
"tf.pywrap_tensorflow",
|
|
"tf.flags",
|
|
"tf.batch_mat_mul_v3",
|
|
"tf.sparse_segment_sum_grad"]
|
|
for path in do_not_document:
|
|
item = tf
|
|
for part in path.split(".")[1:]:
|
|
item = getattr(item, part, None)
|
|
if item is None:
|
|
continue
|
|
doc_controls.do_not_generate_docs(item)
|
|
|
|
base_dirs, code_url_prefixes = base_dir.get_base_dirs_and_prefixes(
|
|
code_url_prefix)
|
|
doc_generator = generate_lib.DocGenerator(
|
|
root_title="TensorFlow 2",
|
|
py_modules=[("tf", tf)],
|
|
base_dir=base_dirs,
|
|
search_hints=search_hints,
|
|
code_url_prefix=code_url_prefixes,
|
|
site_path=site_path,
|
|
visitor_cls=TfExportAwareVisitor,
|
|
private_map=_PRIVATE_MAP,
|
|
extra_docs=_EXTRA_DOCS,
|
|
callbacks=base_dir.get_callbacks())
|
|
|
|
doc_generator.build(output_dir)
|
|
|
|
@contextlib.contextmanager
|
|
def edit_yaml_file(path):
|
|
content = yaml.safe_load(path.read_text())
|
|
yield content
|
|
|
|
with path.open("w") as f:
|
|
yaml.dump(content, f, default_flow_style=False)
|
|
|
|
toc_path = output_dir / "tf/_toc.yaml"
|
|
with edit_yaml_file(toc_path) as toc:
|
|
# Replace the overview path for 'TensorFlow' to
|
|
# `/api_docs/python/tf_overview`. This will be redirected to
|
|
# `/api_docs/python/tf`.
|
|
toc["toc"][0]["section"][0]["path"] = str(site_path / "tf_overview")
|
|
|
|
redirects_path = output_dir / "tf/_redirects.yaml"
|
|
with edit_yaml_file(redirects_path) as redirects:
|
|
redirects["redirects"].append({
|
|
"from": str(site_path / "tf_overview"),
|
|
"to": str(site_path / "tf"),
|
|
})
|
|
|
|
num_files = len(list(output_dir.rglob("*")))
|
|
if num_files < MIN_NUM_FILES_EXPECTED:
|
|
raise ValueError(
|
|
f"The TensorFlow api should be more than {MIN_NUM_FILES_EXPECTED} files"
|
|
f"(found {num_files}).")
|
|
|
|
|
|
def main(argv):
|
|
del argv
|
|
build_docs(
|
|
output_dir=FLAGS.output_dir,
|
|
code_url_prefix=FLAGS.code_url_prefix,
|
|
search_hints=FLAGS.search_hints)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
app.run(main)
|