Improve DefineAndTrace type signature; lingvo 0.12.5 release

PiperOrigin-RevId: 488677433

Author: lingvo-bot

lingvo/core/BUILD

@@ -7,7 +7,7 @@ load(
     "lingvo_proto_py",
     "lingvo_py_binary",
 )
-load("//lingvo:lingvo.bzl", "pytype_library")
+load("//lingvo:lingvo.bzl", "pytype_library", "pytype_strict_library", "pytype_strict_test")
 
 package(
     default_visibility = ["//visibility:public"],
@@ -1414,20 +1414,28 @@ py_library(
     ],
 )
 
-py_test(
+pytype_strict_test(
     name = "py_utils_test",
     size = "medium",
     srcs = ["py_utils_test.py"],
     args = ["--noenable_eager_execution"],
-    deps = [":py_utils_test_lib"],
+    deps = [
+        ":py_utils_test_lib",
+        # Implicit tensorflow functional_ops dependency.
+        # Implicit tensorflow init_ops dependency.
+    ],
 )
 
-py_test(
+pytype_strict_test(
     name = "py_utils_eager_test",
     srcs = ["py_utils_test.py"],
     args = ["--enable_eager_execution"],
     main = "py_utils_test.py",
-    deps = [":py_utils_test_lib"],
+    deps = [
+        ":py_utils_test_lib",
+        # Implicit tensorflow functional_ops dependency.
+        # Implicit tensorflow init_ops dependency.
+    ],
 )
 
 pytype_library(
@@ -1788,19 +1796,20 @@ py_library(
     ],
 )
 
-py_library(
+pytype_strict_library(
     name = "test_utils",
     srcs = ["test_utils.py"],
     deps = [
         ":cluster_factory",
         ":py_utils",
+        ":pytypes",
         "//lingvo:compat",
         # Implicit numpy dependency.
         # Implicit tensorboard/backend/event_processing:event_file_inspector dependency.
     ],
 )
 
-py_test(
+pytype_strict_test(
     name = "test_utils_eager_test",
     srcs = ["test_utils_eager_test.py"],
     args = ["--enable_eager_execution"],
@@ -1810,11 +1819,12 @@ py_test(
     ],
 )
 
-py_test(
+pytype_strict_test(
     name = "test_utils_test",
     srcs = ["test_utils_test.py"],
     args = ["--noenable_eager_execution"],
     deps = [
+        ":py_utils",
         ":test_utils",
         "//lingvo:compat",
     ],
@@ -2600,3 +2610,14 @@ py_strict_test(
         # Implicit numpy dependency.
     ],
 )
+
+pytype_strict_library(
+    name = "pytypes",
+    srcs = ["pytypes.py"],
+    deps = [
+        ":hyperparams",
+        ":nested_map",
+        ":py_utils",
+        # Implicit numpy dependency.
+    ],
+)

lingvo/core/pytypes.py

@@ -0,0 +1,36 @@
+# Copyright 2022 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.
+# ==============================================================================
+"""PyType Typing utilities for lingvo.
+
+A subset of third_party/py/praxis/pytypes.py.
+"""
+
+from typing import List, Tuple, TypeVar, Union, Mapping
+
+from lingvo.core import hyperparams
+from lingvo.core import nested_map
+from lingvo.core import py_utils
+import numpy as np
+
+NpTensor = np.ndarray
+
+NestedMap = nested_map.NestedMap
+Params = hyperparams.Params
+InstantiableParams = hyperparams.InstantiableParams
+
+T = TypeVar('T')
+Nested = Union[T, Tuple[T, ...], List[T], Mapping[str, T], py_utils.NestedMap]
+NestedBool = Nested[bool]
+NestedInt = Nested[int]

lingvo/core/test_utils.py

@@ -26,6 +26,7 @@
 import lingvo.compat as tf
 from lingvo.core import cluster_factory
 from lingvo.core import py_utils
+from lingvo.core import pytypes
 import numpy as np
 
 from tensorboard.backend.event_processing import event_file_inspector
@@ -103,17 +104,40 @@ def gradient(self, ys, xs) -> List[tf.Tensor]:  # pylint: disable=invalid-name
       return tf.gradients(ys=ys, xs=xs)
 
 
+Placeholder = typing.TypeVar('Placeholder', bound=tf.Tensor)
+RetT = typing.TypeVar('RetT')
+
+
+class OverloadedForEager(typing.Protocol, typing.Generic[RetT]):
+
+  def __call__(self, *args: pytypes.Nested[tf.Tensor]) -> RetT:
+    """Needed b/c *args typing support was only added in Py3.11."""
+    ...
+
+
+class OverloadedForGraph(typing.Protocol, typing.Generic[RetT]):
+
+  def __call__(self, *args: pytypes.Nested[Placeholder]) -> RetT:
+    """Needed b/c *args typing support was only added in Py3.11."""
+    ...
+
+
 @typing.overload
 def DefineAndTrace(
-    *tensor_specs_or_placeholders: tf.TensorSpec
-) -> Callable[Callable, Callable]:  # pylint: disable=g-bare-generic
+    *tensor_specs_or_placeholders: pytypes.Nested[tf.TensorSpec]
+) -> Callable[[OverloadedForEager[RetT]],
+              tf.types.experimental.ConcreteFunction]:
+  """Eager case."""
   ...
 
 
 # TODO(jlipschultz): Consider changing the behavior of the graph-mode version of
 # DefineAndTrace to also return a Callable[Callable, Callable] for simplicity.
 @typing.overload
-def DefineAndTrace(*tensor_specs_or_placeholders: tf.Tensor) -> Callable:  # pylint: disable=g-bare-generic
+def DefineAndTrace(
+    *tensor_specs_or_placeholders: pytypes.Nested[Placeholder]
+) -> Callable[[OverloadedForGraph[RetT]], RetT]:
+  """Graph-mode case."""
   ...
 
 

lingvo/lingvo.bzl

@@ -173,7 +173,7 @@ def lingvo_proto_py(name, src, deps = []):
         srcs = [basename + "_pb2.py"],
     )
 
-# Placeholder to use until bazel supports pytype_library.
+# Placeholders to use until bazel supports pytype_{,strict_}{library,test,binary}.
 def pytype_library(name, **kwargs):
     native.py_library(name = name, **kwargs)
 
@@ -183,13 +183,15 @@ def pytype_strict_library(name, **kwargs):
 def pytype_strict_binary(name, **kwargs):
     native.py_binary(name = name, **kwargs)
 
+def py_strict_test(name, **kwargs):
+    native.py_test(name = name, **kwargs)
+
+def pytype_strict_test(name, **kwargs):
+    native.py_test(name = name, **kwargs)
+
 def lingvo_portable_pytype_library(name, deps = [], nonportable_deps = [], **kwargs):
     pytype_library(
         name = name,
         deps = deps + nonportable_deps,
         **kwargs
     )
-
-# Placeholder to use until bazel supports py_strict_test.
-def py_strict_test(name, **kwargs):
-    native.py_test(name = name, **kwargs)

pip_package/setup.py

@@ -20,7 +20,7 @@
 from setuptools.command.install import install
 from setuptools.dist import Distribution
 
-__version__ = '0.12.4'
+__version__ = '0.12.5'
 project_name = 'lingvo'
 if '--project_name' in sys.argv:
   project_name_idx = sys.argv.index('--project_name')