chore: import upstream snapshot with attribution

tensorflow/python/__init__.py

@@ -0,0 +1,34 @@
+# Copyright 2015 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.
+# ==============================================================================
+"""TensorFlow Python init file."""
+
+# Do not add code to //third_party/tensorflow/python/__init__.py.
+# This file is imported whenever TensorFlow is imported.
+# Additional imports in this file could cause the internal
+# import time of TensorFlow to increase by multiple seconds.
+
+
+# Special dunders that we choose to export:
+_exported_dunders = set([
+    '__version__',
+    '__git_version__',
+    '__compiler_version__',
+    '__cxx11_abi_flag__',
+    '__monolithic_build__',
+])
+
+# Expose symbols minus dunders, unless they are allowlisted above.
+# This is necessary to export our dunders.
+__all__ = [s for s in dir() if s in _exported_dunders or not s.startswith('_')]

tensorflow/python/_pywrap_dtensor_device.pyi

@@ -0,0 +1,124 @@
+# 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 typing import ClassVar, Sequence, overload
+
+class Layout:
+    @overload
+    def __init__(self, layout: Layout) -> None: ...
+    @overload
+    def __init__(self, type: LayoutType, sharding_specs: list[str], mesh: Mesh) -> None: ...
+    @overload
+    def __init__(self, layout_proto) -> None: ...
+    @overload
+    def __init__(self, layout_str: str) -> None: ...
+    @overload
+    def __init__(self, mesh: Mesh, rank: int) -> None: ...
+    @overload
+    def __init__(self, mesh: Mesh, rank: int, batch_dim: str, axis: int) -> None: ...
+    @overload
+    def __init__(self, mesh: Mesh) -> None: ...
+    def as_proto(self, *args, **kwargs): ...
+    def global_shape_from_local_shape(self, local_shape: list[int]) -> tuple: ...
+    def is_batch_parallel(self) -> bool: ...
+    def is_fully_replicated(self) -> bool: ...
+    def is_single_device(self) -> bool: ...
+    def local_shape_from_global_shape(self, global_shape: list[int]) -> tuple: ...
+    def num_shards(self, idx: int) -> int: ...
+    def to_parted(self) -> Layout: ...
+    def to_string(self) -> str: ...
+    def __eq__(self, arg0: Layout) -> bool: ...
+    @property
+    def mesh(self) -> Mesh: ...
+    @property
+    def rank(self) -> int: ...
+    @property
+    def sharding_specs(self) -> list[str]: ...
+    @property
+    def type(self) -> LayoutType: ...
+
+class LayoutType:
+    __members__: ClassVar[dict] = ...  # read-only
+    PARTED: ClassVar[LayoutType] = ...
+    SINGLE_DEVICE: ClassVar[LayoutType] = ...
+    STATIC: ClassVar[LayoutType] = ...
+    __entries: ClassVar[dict] = ...
+    def __init__(self, value: int) -> None: ...
+    def __eq__(self, other: object) -> bool: ...
+    def __hash__(self) -> int: ...
+    def __index__(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __ne__(self, other: object) -> bool: ...
+    @property
+    def name(self) -> str: ...
+    @property
+    def value(self) -> int: ...
+
+class Mesh:
+    @overload
+    def __init__(self, mesh: Mesh) -> None: ...
+    @overload
+    def __init__(self, arg0: str, arg1: list[str], arg2: list[int], arg3: list[int], arg4: list[str], arg5: list[int], arg6: list[str], arg7: bool) -> None: ...
+    @overload
+    def __init__(self, single_device: str) -> None: ...
+    @overload
+    def __init__(self, mesh_proto) -> None: ...
+    @overload
+    def __init__(self, mesh_str: str) -> None: ...
+    def as_proto(self, *args, **kwargs): ...
+    def contains_dim(self, dim_name: str) -> bool: ...
+    def device_location(self, arg0: int) -> list[int]: ...
+    def device_type(self) -> str: ...
+    def dim_size(self, dim_name: str) -> int: ...
+    def global_device_ids(self) -> Sequence[int]: ...
+    def global_devices(self) -> list[str]: ...
+    def host_mesh(self) -> Mesh: ...
+    def is_remote(self) -> bool: ...
+    def is_single_device(self) -> bool: ...
+    def local_device_ids(self) -> Sequence[int]: ...
+    def local_devices(self) -> Sequence[str]: ...
+    def min_global_device_id(self) -> int: ...
+    def num_local_devices(self) -> int: ...
+    def shape(self) -> list[int]: ...
+    def to_string(self) -> str: ...
+    def use_xla_spmd(self) -> bool: ...
+    def __contains__(self, dim_name: str) -> bool: ...
+    def __eq__(self, arg0: Mesh) -> bool: ...
+    @property
+    def dim_names(self) -> list[str]: ...
+    @property
+    def name(self) -> str: ...
+    @property
+    def single_device(self) -> str: ...
+    @property
+    def size(self) -> int: ...
+
+def AddMesh(arg0, arg1: str, arg2: bool) -> None: ...
+def Allocate(arg0: str, arg1: bool, arg2: int) -> object: ...
+def ClearTPUCoreIDs(arg0) -> None: ...
+def ExperimentalClearDefaultLayout(arg0) -> None: ...
+def ExperimentalClearDefaultMesh(arg0) -> None: ...
+def ExperimentalSetDefaultLayout(arg0, arg1: str) -> None: ...
+def ExperimentalSetDefaultMesh(arg0, arg1: str) -> None: ...
+def FetchLayout(arg0: object, arg1: object, arg2) -> object: ...
+def GetStats(arg0: object, arg1) -> dict[str, int]: ...
+def IsDTensor(arg0: object, arg1: object, arg2) -> bool: ...
+def IsSparseDTensor(arg0: object, arg1: object, arg2) -> bool: ...
+def Pack(arg0: object, arg1: object, arg2: str, arg3, arg4: bool) -> object: ...
+def SetIteratorElementLayouts(arg0: object, arg1: object, arg2: list[str], arg3) -> None: ...
+def SetTPUCoreIDs(arg0, arg1: str, arg2: list[int]) -> None: ...
+def TPUCoreIDsToLocations(arg0: object, arg1, arg2: list[int]) -> list[list[int]]: ...
+def TPUCoreLocationsToIDs(arg0: object, arg1, arg2: list[list[int]]) -> list[int]: ...
+def Unpack(arg0: object, arg1: object, arg2) -> object: ...

tensorflow/python/_pywrap_mlir.pyi

@@ -0,0 +1,27 @@
+# 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 typing import overload
+
+def ExperimentalConvertSavedModelToMlir(arg0: str, arg1: str, arg2: bool) -> str: ...
+def ExperimentalConvertSavedModelV1ToMlir(arg0: str, arg1: str, arg2: str, arg3: bool, arg4: bool, arg5: bool, arg6: bool) -> str: ...
+def ExperimentalConvertSavedModelV1ToMlirLite(arg0: str, arg1: str, arg2: str, arg3: bool, arg4: bool) -> str: ...
+def ExperimentalRunPassPipeline(arg0: str, arg1: str, arg2: bool) -> str: ...
+def ExperimentalWriteBytecode(arg0: str, arg1: str) -> None: ...
+def ImportFunction(arg0: object, arg1: str, arg2: str, arg3: bool) -> str: ...
+@overload
+def ImportGraphDef(arg0: str, arg1: str, arg2: bool) -> str: ...
+@overload
+def ImportGraphDef(arg0: str, arg1: str, arg2: bool, arg3: str, arg4: str, arg5: str, arg6: str) -> str: ...

tensorflow/python/_pywrap_parallel_device.pyi

@@ -0,0 +1,16 @@
+# 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.
+# ==============================================================================
+
+def GetParallelDeviceCapsules(arg0: str, arg1: list[str]) -> object: ...

tensorflow/python/_pywrap_py_exception_registry.pyi

@@ -0,0 +1,62 @@
+# 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 typing import ClassVar
+
+TF_ABORTED: TF_Code
+TF_CANCELLED: TF_Code
+TF_DATA_LOSS: TF_Code
+TF_DEADLINE_EXCEEDED: TF_Code
+TF_FAILED_PRECONDITION: TF_Code
+TF_INTERNAL: TF_Code
+TF_INVALID_ARGUMENT: TF_Code
+TF_OK: TF_Code
+TF_OUT_OF_RANGE: TF_Code
+TF_PERMISSION_DENIED: TF_Code
+TF_RESOURCE_EXHAUSTED: TF_Code
+TF_UNAUTHENTICATED: TF_Code
+TF_UNIMPLEMENTED: TF_Code
+TF_UNKNOWN: TF_Code
+
+class TF_Code:
+    __members__: ClassVar[dict] = ...  # read-only
+    TF_ABORTED: ClassVar[TF_Code] = ...
+    TF_CANCELLED: ClassVar[TF_Code] = ...
+    TF_DATA_LOSS: ClassVar[TF_Code] = ...
+    TF_DEADLINE_EXCEEDED: ClassVar[TF_Code] = ...
+    TF_FAILED_PRECONDITION: ClassVar[TF_Code] = ...
+    TF_INTERNAL: ClassVar[TF_Code] = ...
+    TF_INVALID_ARGUMENT: ClassVar[TF_Code] = ...
+    TF_OK: ClassVar[TF_Code] = ...
+    TF_OUT_OF_RANGE: ClassVar[TF_Code] = ...
+    TF_PERMISSION_DENIED: ClassVar[TF_Code] = ...
+    TF_RESOURCE_EXHAUSTED: ClassVar[TF_Code] = ...
+    TF_UNAUTHENTICATED: ClassVar[TF_Code] = ...
+    TF_UNIMPLEMENTED: ClassVar[TF_Code] = ...
+    TF_UNKNOWN: ClassVar[TF_Code] = ...
+    __entries: ClassVar[dict] = ...
+    def __init__(self, value: int) -> None: ...
+    def __eq__(self, other: object) -> bool: ...
+    def __hash__(self) -> int: ...
+    def __index__(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __ne__(self, other: object) -> bool: ...
+    @property
+    def name(self) -> str: ...
+    @property
+    def value(self) -> int: ...
+
+def PyExceptionRegistry_Init(arg0: object) -> None: ...
+def PyExceptionRegistry_Lookup(arg0: TF_Code) -> None: ...

tensorflow/python/_pywrap_quantize_training.pyi

@@ -0,0 +1,16 @@
+# 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.
+# ==============================================================================
+
+def DoQuantizeTrainingOnGraphDefHelper(arg0: object, arg1: int) -> object: ...

tensorflow/python/_pywrap_sanitizers.pyi

@@ -0,0 +1,19 @@
+# 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.
+# ==============================================================================
+
+def is_asan_enabled() -> bool: ...
+def is_msan_enabled() -> bool: ...
+def is_tsan_enabled() -> bool: ...
+def is_ubsan_enabled() -> bool: ...

tensorflow/python/_pywrap_tensorflow.def

@@ -0,0 +1,868 @@
+; Copyright 2024 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.
+
+; This file is manually created by removing all entries and running the Windows
+; presubmit (tensorflow/official/presubmit/critique/windows_x86_2022).
+; The linker then fails with `undefined symbol` errors, and the missing symbols
+; are added, sorted and de-duplicated here.
+; NOTE: Symbols beginning with `__imp_` should have that prefix removed, e.g.
+; `__imp_??1OpDef@tensorflow@@UEAA@XZ` becomes `??1OpDef@tensorflow@@UEAA@XZ`.
+
+; go/keep-sorted start skip_lines=1
+EXPORTS
+ ??$Accept@PEAVHloInstruction@xla@@@HloInstruction@xla@@QEAA?AVStatus@lts_20260526@absl@@PEAV?$DfsHloVisitorBase@PEAVHloInstruction@xla@@@1@_N11@Z
+ ??$Accept@PEBVHloInstruction@xla@@@HloInstruction@xla@@QEAA?AVStatus@lts_20260526@absl@@PEAV?$DfsHloVisitorBase@PEBVHloInstruction@xla@@@1@_N11@Z
+ ??$AllocateBackingArray@$07V?$allocator@D@std@@@container_internal@lts_20260526@absl@@YAPEAXPEAX_K@Z
+ ??$CopyToEncodedBuffer@$0A@@LogMessage@log_internal@lts_20260526@absl@@AEAAXV?$basic_string_view@DU?$char_traits@D@std@@@std@@@Z
+ ??$DeallocateBackingArray@$07V?$allocator@D@std@@@container_internal@lts_20260526@absl@@YAXPEAX_KPEAW4ctrl_t@012@11_N@Z
+ ??$Dispatch@H@FormatArgImpl@str_format_internal@lts_20260526@absl@@CA_NTData@0123@VFormatConversionSpecImpl@123@PEAX@Z
+ ??$Dispatch@PEBD@FormatArgImpl@str_format_internal@lts_20260526@absl@@CA_NTData@0123@VFormatConversionSpecImpl@123@PEAX@Z
+ ??$Dispatch@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@FormatArgImpl@str_format_internal@lts_20260526@absl@@CA_NTData@0123@VFormatConversionSpecImpl@123@PEAX@Z
+ ??$Dispatch@V?$basic_string_view@DU?$char_traits@D@std@@@std@@@FormatArgImpl@str_format_internal@lts_20260526@absl@@CA_NTData@0123@VFormatConversionSpecImpl@123@PEAX@Z
+ ??$GrowSooTableToNextCapacityAndPrepareInsert@$07$00@container_internal@lts_20260526@absl@@YA_KAEAVCommonFields@012@AEBUPolicyFunctions@012@V?$FunctionRef@$$A6A_K_K@Z@12@_N@Z
+ ??$GrowSooTableToNextCapacityAndPrepareInsert@$0A@$0A@@container_internal@lts_20260526@absl@@YA_KAEAVCommonFields@012@AEBUPolicyFunctions@012@V?$FunctionRef@$$A6A_K_K@Z@12@_N@Z
+ ??$GrowSooTableToNextCapacityAndPrepareInsert@$0BA@$00@container_internal@lts_20260526@absl@@YA_KAEAVCommonFields@012@AEBUPolicyFunctions@012@V?$FunctionRef@$$A6A_K_K@Z@12@_N@Z
+ ??$MakeErrorImpl@$02@status_internal@lts_20260526@absl@@YA?AVStatus@12@V?$basic_string_view@DU?$char_traits@D@std@@@std@@VSourceLocation@12@@Z
+ ??$MakeErrorImpl@$04@status_internal@lts_20260526@absl@@YA?AVStatus@12@V?$basic_string_view@DU?$char_traits@D@std@@@std@@VSourceLocation@12@@Z
+ ??$MakeErrorImpl@$05@status_internal@lts_20260526@absl@@YA?AVStatus@12@V?$basic_string_view@DU?$char_traits@D@std@@@std@@VSourceLocation@12@@Z
+ ??$MakeErrorImpl@$06@status_internal@lts_20260526@absl@@YA?AVStatus@12@V?$basic_string_view@DU?$char_traits@D@std@@@std@@VSourceLocation@12@@Z
+ ??$MakeErrorImpl@$08@status_internal@lts_20260526@absl@@YA?AVStatus@12@V?$basic_string_view@DU?$char_traits@D@std@@@std@@VSourceLocation@12@@Z
+ ??$MakeErrorImpl@$0M@@status_internal@lts_20260526@absl@@YA?AVStatus@12@V?$basic_string_view@DU?$char_traits@D@std@@@std@@VSourceLocation@12@@Z
+ ??$MakeErrorImpl@$0N@@status_internal@lts_20260526@absl@@YA?AVStatus@12@V?$basic_string_view@DU?$char_traits@D@std@@@std@@VSourceLocation@12@@Z
+ ??$Visit@PEAVHloInstruction@xla@@@HloInstruction@xla@@QEAA?AVStatus@lts_20260526@absl@@PEAV?$DfsHloVisitorBase@PEAVHloInstruction@xla@@@1@@Z
+ ??$Visit@PEBVHloInstruction@xla@@@HloInstruction@xla@@QEAA?AVStatus@lts_20260526@absl@@PEAV?$DfsHloVisitorBase@PEBVHloInstruction@xla@@@1@@Z
+ ??$eigen_device@UThreadPoolDevice@Eigen@@@OpKernelContext@tensorflow@@QEBAAEBUThreadPoolDevice@Eigen@@XZ
+ ??0?$TensorShapeBase@VPartialTensorShape@tensorflow@@@tensorflow@@QEAA@XZ
+ ??0?$TensorShapeBase@VTensorShape@tensorflow@@@tensorflow@@QEAA@V?$Span@$$CB_J@lts_20260526@absl@@@Z
+ ??0?$TensorShapeBase@VTensorShape@tensorflow@@@tensorflow@@QEAA@XZ
+ ??0ArithDialect@arith@mlir@@AEAA@PEAVMLIRContext@2@@Z
+ ??0AttrSlice@tensorflow@@QEAA@AEBVNodeDef@1@@Z
+ ??0AttrValue@tensorflow@@IEAA@PEAVArena@protobuf@google@@@Z
+ ??0BufferedInputStream@io@tsl@@QEAA@PEAVInputStreamInterface@12@_K_N@Z
+ ??0CalibrationOptions@quantization@stablehlo@@IEAA@PEAVArena@protobuf@google@@@Z
+ ??0CalibrationStatistics@calibrator@tensorflow@@IEAA@PEAVArena@protobuf@google@@@Z
+ ??0CancellationManager@tsl@@QEAA@XZ
+ ??0CheckpointReader@checkpoint@tensorflow@@QEAA@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@PEAUTSL_Status@@@Z
+ ??0ConfigProto@tensorflow@@IEAA@PEAVArena@protobuf@google@@@Z
+ ??0ConfigProto@tensorflow@@IEAA@PEAVArena@protobuf@google@@AEBV01@@Z
+ ??0CoordinatedTask@tensorflow@@IEAA@PEAVArena@protobuf@google@@@Z
+ ??0Cord@lts_20260526@absl@@AEAA@V?$basic_string_view@DU?$char_traits@D@std@@@std@@W4MethodIdentifier@CordzUpdateTracker@cord_internal@12@@Z
+ ??0CostAnalyzer@grappler@tensorflow@@QEAA@AEBUGrapplerItem@12@PEAVCluster@12@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@Z
+ ??0DataServiceMetadata@data@tensorflow@@IEAA@PEAVArena@protobuf@google@@@Z
+ ??0DeviceAttributes@tensorflow@@IEAA@PEAVArena@protobuf@google@@AEBV01@@Z
+ ??0DeviceProperties@tensorflow@@IEAA@PEAVArena@protobuf@google@@@Z
+ ??0DialectRegistry@mlir@@QEAA@XZ
+ ??0DispatcherConfig@experimental@data@tensorflow@@IEAA@PEAVArena@protobuf@google@@@Z
+ ??0EventsWriter@tensorflow@@QEAA@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@Z
+ ??0ExportedModel@quantization@tensorflow@@IEAA@PEAVArena@protobuf@google@@@Z
+ ??0FileCheck@llvm@@QEAA@UFileCheckRequest@1@@Z
+ ??0FingerprintDef@tensorflow@@IEAA@PEAVArena@protobuf@google@@@Z
+ ??0FingerprintDef@tensorflow@@IEAA@PEAVArena@protobuf@google@@AEBV01@@Z
+ ??0FrozenStackTrace@tensorflow@@QEAA@V?$Span@$$CBUStackFrame@tsl@@@lts_20260526@absl@@0@Z
+ ??0FuncDialect@func@mlir@@AEAA@PEAVMLIRContext@2@@Z
+ ??0FunctionDef@tensorflow@@IEAA@PEAVArena@protobuf@google@@@Z
+ ??0FunctionParameterCanonicalizer@tensorflow@@QEAA@V?$Span@PEBD@lts_20260526@absl@@V?$Span@PEAU_object@@@34@@Z
+ ??0FunctionRecord@tensorflow@@QEAA@$$QEAVFunctionDef@1@$$QEAV?$flat_hash_map@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$shared_ptr@VAbstractStackTrace@tensorflow@@@2@UStringHash@container_internal@lts_20260526@absl@@UStringEq@567@V?$allocator@U?$pair@$$CBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$shared_ptr@VAbstractStackTrace@tensorflow@@@2@@std@@@2@@lts_20260526@absl@@_N@Z
+ ??0GraphDebugInfoBuilder@tensorflow@@QEAA@XZ
+ ??0GraphDef@tensorflow@@IEAA@PEAVArena@protobuf@google@@@Z
+ ??0GraphDef@tensorflow@@IEAA@PEAVArena@protobuf@google@@AEBV01@@Z
+ ??0HistogramProto@tensorflow@@IEAA@PEAVArena@protobuf@google@@@Z
+ ??0LayoutProto@dtensor@tensorflow@@IEAA@PEAVArena@protobuf@google@@@Z
+ ??0LayoutProto@dtensor@tensorflow@@IEAA@PEAVArena@protobuf@google@@AEBV012@@Z
+ ??0LogMessage@log_internal@lts_20260526@absl@@QEAA@PEBDHUInfoTag@0123@@Z
+ ??0LogMessageFatal@log_internal@lts_20260526@absl@@QEAA@PEBDH0@Z
+ ??0LogMessageFatal@log_internal@lts_20260526@absl@@QEAA@PEBDH@Z
+ ??0MLIRContext@mlir@@QEAA@AEBVDialectRegistry@1@W4Threading@01@@Z
+ ??0MakeErrorStream@status_macros@xla@@QEAA@PEBDHW4StatusCode@lts_20260526@absl@@@Z
+ ??0MeasuringCostEstimator@grappler@tensorflow@@QEAA@PEAVCluster@12@HH@Z
+ ??0MeshProto@dtensor@tensorflow@@IEAA@PEAVArena@protobuf@google@@@Z
+ ??0MeshProto@dtensor@tensorflow@@IEAA@PEAVArena@protobuf@google@@AEBV012@@Z
+ ??0MessageDifferencer@util@protobuf@google@@QEAA@XZ
+ ??0MetaGraphDef@tensorflow@@IEAA@PEAVArena@protobuf@google@@@Z
+ ??0MetaOptimizer@grappler@tensorflow@@QEAA@PEAVDeviceBase@2@AEBVConfigProto@2@@Z
+ ??0ModelAnalyzer@grappler@tensorflow@@QEAA@AEBUGrapplerItem@12@@Z
+ ??0Name@register_kernel@tensorflow@@QEAA@PEBD@Z
+ ??0NamedDevice@tensorflow@@IEAA@PEAVArena@protobuf@google@@@Z
+ ??0NamedDevice@tensorflow@@IEAA@PEAVArena@protobuf@google@@AEBV01@@Z
+ ??0NodeDef@tensorflow@@IEAA@PEAVArena@protobuf@google@@AEBV01@@Z
+ ??0OpDef@tensorflow@@IEAA@PEAVArena@protobuf@google@@AEBV01@@Z
+ ??0OpDefBuilder@tensorflow@@QEAA@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@Z
+ ??0OpKernel@tensorflow@@QEAA@PEAVOpKernelConstruction@1@@Z
+ ??0OpLevelCostEstimator@grappler@tensorflow@@QEAA@XZ
+ ??0OpPerformanceList@tensorflow@@IEAA@PEAVArena@protobuf@google@@@Z
+ ??0OstreamView@LogMessage@log_internal@lts_20260526@absl@@QEAA@AEAULogMessageData@1234@@Z
+ ??0PyInstanceChecker@py_dispatch@tensorflow@@QEAA@AEBV?$vector@PEAU_object@@V?$allocator@PEAU_object@@@std@@@std@@@Z
+ ??0PySignatureChecker@py_dispatch@tensorflow@@QEAA@V?$vector@U?$pair@HV?$shared_ptr@VPyTypeChecker@py_dispatch@tensorflow@@@std@@@std@@V?$allocator@U?$pair@HV?$shared_ptr@VPyTypeChecker@py_dispatch@tensorflow@@@std@@@std@@@2@@std@@@Z
+ ??0PythonAPIDispatcher@py_dispatch@tensorflow@@QEAA@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$Span@PEBD@lts_20260526@absl@@V?$Span@PEAU_object@@@67@@Z
+ ??0PythonAPIInfo@tensorflow@@QEAA@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@Z
+ ??0QuantizationConfig@quantization@stablehlo@@IEAA@PEAVArena@protobuf@google@@@Z
+ ??0QuantizationOptions@quantization@tensorflow@@IEAA@PEAVArena@protobuf@google@@@Z
+ ??0QueueRunnerDef@tensorflow@@IEAA@PEAVArena@protobuf@google@@AEBV01@@Z
+ ??0RandomAccessInputStream@io@tsl@@QEAA@PEAVRandomAccessFile@2@_N@Z
+ ??0RecordReader@io@tsl@@QEAA@PEAVRandomAccessFile@2@AEBURecordReaderOptions@12@@Z
+ ??0RecordWriter@io@tsl@@QEAA@PEAVWritableFile@2@AEBURecordWriterOptions@12@@Z
+ ??0RepresentativeDatasetFile@quantization@tensorflow@@IEAA@PEAVArena@protobuf@google@@@Z
+ ??0RunMetadata@tensorflow@@IEAA@PEAVArena@protobuf@google@@@Z
+ ??0SCFDialect@scf@mlir@@AEAA@PEAVMLIRContext@2@@Z
+ ??0SessionOptions@tensorflow@@QEAA@XZ
+ ??0ShapeDialect@shape@mlir@@AEAA@PEAVMLIRContext@2@@Z
+ ??0SignatureDef@tensorflow@@IEAA@PEAVArena@protobuf@google@@@Z
+ ??0SimpleFieldComparator@util@protobuf@google@@QEAA@XZ
+ ??0SingleMachine@grappler@tensorflow@@QEAA@HHH@Z
+ ??0SourceMgr@llvm@@QEAA@XZ
+ ??0SourceMgrDiagnosticHandler@mlir@@QEAA@AEAVSourceMgr@llvm@@PEAVMLIRContext@1@$$QEAV?$unique_function@$$A6A_NVLocation@mlir@@@Z@3@@Z
+ ??0SparseCoreLayoutStacker@tpu@tensorflow@@QEAA@H_NH@Z
+ ??0SrcBuffer@SourceMgr@llvm@@QEAA@$$QEAU012@@Z
+ ??0StatSummarizer@tensorflow@@QEAA@AEBUStatSummarizerOptions@tsl@@@Z
+ ??0StatSummarizer@tensorflow@@QEAA@AEBVGraphDef@1@@Z
+ ??0StepStats@tensorflow@@IEAA@PEAVArena@protobuf@google@@@Z
+ ??0TFRDialect@TFR@mlir@@QEAA@PEAVMLIRContext@2@@Z
+ ??0TapeContext@gradients@tensorflow@@QEAA@PEAVAbstractContext@2@PEAVTape@12@AEBVGradientRegistry@12@@Z
+ ??0Tensor@tensorflow@@QEAA@XZ
+ ??0TensorFlowDialect@TF@mlir@@QEAA@PEAVMLIRContext@2@@Z
+ ??0VirtualCluster@grappler@tensorflow@@QEAA@AEBV?$unordered_map@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@VDeviceProperties@tensorflow@@U?$hash@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@U?$equal_to@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@V?$allocator@U?$pair@$$CBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@VDeviceProperties@tensorflow@@@std@@@2@@std@@@Z
+ ??0WorkerConfig@experimental@data@tensorflow@@IEAA@PEAVArena@protobuf@google@@@Z
+ ??1ApiDefMap@tensorflow@@QEAA@XZ
+ ??1AttrValue@tensorflow@@UEAA@XZ
+ ??1Block@mlir@@QEAA@XZ
+ ??1BundleReader@tensorflow@@QEAA@XZ
+ ??1CalibrationOptions@quantization@stablehlo@@UEAA@XZ
+ ??1CalibrationStatistics@calibrator@tensorflow@@UEAA@XZ
+ ??1CancellationManager@tsl@@QEAA@XZ
+ ??1ConfigProto@tensorflow@@UEAA@XZ
+ ??1CoordinatedTask@tensorflow@@UEAA@XZ
+ ??1DataServiceMetadata@data@tensorflow@@UEAA@XZ
+ ??1DeviceAttributes@tensorflow@@UEAA@XZ
+ ??1DeviceProperties@tensorflow@@UEAA@XZ
+ ??1DispatcherConfig@experimental@data@tensorflow@@UEAA@XZ
+ ??1EagerExecutor@tensorflow@@QEAA@XZ
+ ??1EventsWriter@tensorflow@@QEAA@XZ
+ ??1ExportedModel@quantization@tensorflow@@UEAA@XZ
+ ??1FileCheck@llvm@@QEAA@XZ
+ ??1FingerprintDef@tensorflow@@UEAA@XZ
+ ??1FunctionDef@tensorflow@@UEAA@XZ
+ ??1GraphDef@tensorflow@@UEAA@XZ
+ ??1HistogramProto@tensorflow@@UEAA@XZ
+ ??1Impl@MakeErrorStream@status_macros@xla@@QEAA@XZ
+ ??1KernelDefBuilder@tensorflow@@QEAA@XZ
+ ??1LayoutProto@dtensor@tensorflow@@UEAA@XZ
+ ??1LogMessage@log_internal@lts_20260526@absl@@QEAA@XZ
+ ??1LogMessageFatal@log_internal@lts_20260526@absl@@QEAA@XZ
+ ??1MLIRContext@mlir@@QEAA@XZ
+ ??1MeshProto@dtensor@tensorflow@@UEAA@XZ
+ ??1MessageDifferencer@util@protobuf@google@@QEAA@XZ
+ ??1MetaGraphDef@tensorflow@@UEAA@XZ
+ ??1NamedDevice@tensorflow@@UEAA@XZ
+ ??1NodeDef@tensorflow@@UEAA@XZ
+ ??1OpDef@tensorflow@@UEAA@XZ
+ ??1OpInfo_TensorProperties@tensorflow@@UEAA@XZ
+ ??1OpKernel@tensorflow@@UEAA@XZ
+ ??1OpPerformanceList@tensorflow@@UEAA@XZ
+ ??1OstreamView@LogMessage@log_internal@lts_20260526@absl@@UEAA@XZ
+ ??1ProfilerServer@profiler@tsl@@QEAA@XZ
+ ??1ProfilerSession@tsl@@QEAA@XZ
+ ??1PyContextManager@tensorflow@@QEAA@XZ
+ ??1QuantizationConfig@quantization@stablehlo@@UEAA@XZ
+ ??1QuantizationOptions@quantization@tensorflow@@UEAA@XZ
+ ??1QueueRunnerDef@tensorflow@@UEAA@XZ
+ ??1RecordWriter@io@tsl@@QEAA@XZ
+ ??1RepresentativeDatasetFile@quantization@tensorflow@@UEAA@XZ
+ ??1ResourceHandle@tensorflow@@QEAA@XZ
+ ??1RunMetadata@tensorflow@@UEAA@XZ
+ ??1SignatureDef@tensorflow@@UEAA@XZ
+ ??1SimpleFieldComparator@util@protobuf@google@@UEAA@XZ
+ ??1SourceMgr@llvm@@QEAA@XZ
+ ??1SourceMgrDiagnosticHandler@mlir@@QEAA@XZ
+ ??1SparseCoreTableLayout@tpu@tensorflow@@UEAA@XZ
+ ??1SparseCoreTableLayouts@tpu@tensorflow@@UEAA@XZ
+ ??1SrcBuffer@SourceMgr@llvm@@QEAA@XZ
+ ??1StatSummarizer@tensorflow@@QEAA@XZ
+ ??1StepStats@tensorflow@@UEAA@XZ
+ ??1TapeTensor@gradients@tensorflow@@QEAA@XZ
+ ??1Tensor@tensorflow@@QEAA@XZ
+ ??1ThreadPool@thread@tsl@@QEAA@XZ
+ ??1Variant@tensorflow@@QEAA@XZ
+ ??1WorkerConfig@experimental@data@tensorflow@@UEAA@XZ
+ ??6LogMessage@log_internal@lts_20260526@absl@@QEAAAEAV0123@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@Z
+ ??8Layout@dtensor@tensorflow@@QEBA_NAEBV012@@Z
+ ??8Mesh@dtensor@tensorflow@@QEBA_NAEBV012@@Z
+ ??BCord@lts_20260526@absl@@QEBA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@XZ
+ ??ROpDefBuilderWrapper@register_op@tensorflow@@QEAA?AUInitOnStartupMarker@2@XZ
+ ?AccumulateStackTrace@GraphDebugInfoBuilder@tensorflow@@QEAAXV?$shared_ptr@VAbstractStackTrace@tensorflow@@@std@@V?$basic_string_view@DU?$char_traits@D@std@@@4@AEBUOptions@12@@Z
+ ?Add@InferenceContext@shape_inference@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@VDimensionHandle@23@UDimensionOrConstant@23@PEAV723@@Z
+ ?Add@ThreadSafeHistogram@histogram@tsl@@QEAAXN@Z
+ ?AddControlEdge@Graph@tensorflow@@QEAAPEBVEdge@2@PEAVNode@2@0_N@Z
+ ?AddDevices@DeviceFactory@tensorflow@@SA?AVStatus@lts_20260526@absl@@AEBUSessionOptions@2@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@PEAV?$vector@V?$unique_ptr@VDevice@tensorflow@@U?$default_delete@VDevice@tensorflow@@@std@@@std@@V?$allocator@V?$unique_ptr@VDevice@tensorflow@@U?$default_delete@VDevice@tensorflow@@@std@@@std@@@2@@8@@Z
+ ?AddMesh@dtensor@tensorflow@@YAXAEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@PEAX_NPEAUTSL_Status@@@Z
+ ?AddRegisterer@gradients@tensorflow@@YAPEAVGradientFunction@12@AEBUForwardOperation@12@@Z
+ ?AddStep@tfprof@tensorflow@@YAN_JPEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@11@Z
+ ?AddTable@SparseCoreLayoutStacker@tpu@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@V?$basic_string_view@DU?$char_traits@D@std@@@std@@_J1011@Z
+ ?AddV2@ops@tensorflow@@YA?AVStatus@lts_20260526@absl@@PEAVAbstractContext@2@QEAVAbstractTensorHandle@2@1PEAPEAV72@PEBD3@Z
+ ?AddWhileInputHack@tensorflow@@YAXPEAUTF_Graph@@UTF_Output@@PEAUTF_Operation@@PEAUTSL_Status@@@Z
+ ?AllocateDTensorDevice@dtensor@tensorflow@@YAXV?$basic_string_view@DU?$char_traits@D@std@@@std@@PEAUTFE_CustomDevice@@PEAPEAX_NHPEAUTSL_Status@@@Z
+ ?AllowSoftPlacement@Cluster@grappler@tensorflow@@QEAAX_N@Z
+ ?AppendGraphDebugInfoStr@GraphDebugInfoBuilder@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@V?$basic_string_view@DU?$char_traits@D@std@@@std@@0@Z
+ ?AppendPieces@strings_internal@lts_20260526@absl@@YAXPEAV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$initializer_list@V?$basic_string_view@DU?$char_traits@D@std@@@std@@@5@@Z
+ ?ApplyAllOptionOverrides@CompileOptions@xla@@QEAA?AVStatus@lts_20260526@absl@@XZ
+ ?AreCompatibleDevNames@DeviceNameUtils@tsl@@SA_NAEBUParsedName@12@0@Z
+ ?AssertSameStructure@swig@tensorflow@@YAPEAU_object@@PEAU3@0_N1@Z
+ ?AssertSameStructureForData@swig@tensorflow@@YAPEAU_object@@PEAU3@0_N@Z
+ ?AsyncCheckpointWriteDuration@metrics@tensorflow@@YAAEAVSamplerCell@monitoring@tsl@@V?$basic_string_view@DU?$char_traits@D@std@@@std@@@Z
+ ?Attr@OpDefBuilder@tensorflow@@QEAAAEAV12@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@Z
+ ?AttrValueToPyObject@tensorflow@@YA?AV?$unique_ptr@U_object@@UPyDecrefDeleter@detail@tensorflow@@@std@@AEBVAttrValue@1@@Z
+ ?AttributeTypeFromName@tensorflow@@YA?AW4AttributeType@1@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@Z
+ ?BatchShardedOnMesh@Layout@dtensor@tensorflow@@SA?AV123@AEBVMesh@23@HAEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@H@Z
+ ?BoundPort@GrpcDataServerBase@data@tensorflow@@QEAAHXZ
+ ?Build@KernelDefBuilder@tensorflow@@QEAAPEBVKernelDef@2@XZ
+ ?Canonicalize@FunctionParameterCanonicalizer@tensorflow@@QEAA_NPEAU_object@@0V?$Span@PEAU_object@@@lts_20260526@absl@@@Z
+ ?Capture@StackTrace@tensorflow@@SA?AV?$shared_ptr@VStackTrace@tensorflow@@@std@@H@Z
+ ?CatPieces@strings_internal@lts_20260526@absl@@YA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$initializer_list@V?$basic_string_view@DU?$char_traits@D@std@@@std@@@5@@Z
+ ?Check@PyListChecker@py_dispatch@tensorflow@@UEAA?AW4MatchType@PyTypeChecker@23@PEAU_object@@@Z
+ ?Check@PyUnionChecker@py_dispatch@tensorflow@@UEAA?AW4MatchType@PyTypeChecker@23@PEAU_object@@@Z
+ ?CheckCanonicalizedArgs@PySignatureChecker@py_dispatch@tensorflow@@QEBA_NV?$Span@PEAU_object@@@lts_20260526@absl@@@Z
+ ?CheckDimsEqual@TensorShape@tensorflow@@AEBAXH@Z
+ ?CheckIsAlignedAndSingleElement@Tensor@tensorflow@@AEBAXXZ
+ ?CheckNotDone@MakeErrorStream@status_macros@xla@@AEBAXXZ
+ ?CheckNotInComputeAsync@tensorflow@@YAXPEAVOpKernelContext@1@PEBD@Z
+ ?CheckOpDeprecation@tensorflow@@YA?AVStatus@lts_20260526@absl@@AEBVOpDef@1@H@Z
+ ?CheckTypeAndIsAligned@Tensor@tensorflow@@AEBAXW4DataType@2@@Z
+ ?CheckpointReadDuration@metrics@tensorflow@@YAAEAVSamplerCell@monitoring@tsl@@V?$basic_string_view@DU?$char_traits@D@std@@@std@@@Z
+ ?CheckpointSize@metrics@tensorflow@@YAAEAVCounterCell@monitoring@tsl@@V?$basic_string_view@DU?$char_traits@D@std@@@std@@_J@Z
+ ?CheckpointWriteDuration@metrics@tensorflow@@YAAEAVSamplerCell@monitoring@tsl@@V?$basic_string_view@DU?$char_traits@D@std@@@std@@@Z
+ ?Clear@TFE_TensorHandleCache@tensorflow@@QEAAXXZ
+ ?ClearAttr@tensorflow@@YAXPEAUTF_Graph@@PEAUTF_Operation@@PEBDPEAUTSL_Status@@@Z
+ ?ClearBackingArray@container_internal@lts_20260526@absl@@YAXAEAVCommonFields@123@AEBUPolicyFunctions@123@PEAX_N@Z
+ ?ClearTPUCoreIDs@dtensor@tensorflow@@YAXPEAX@Z
+ ?Close@DebugEventsWriter@tfdbg@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@XZ
+ ?Close@EventsWriter@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@XZ
+ ?Close@RecordWriter@io@tsl@@QEAA?AVStatus@lts_20260526@absl@@XZ
+ ?CollectUnreachableRoots@HloComputation@xla@@AEBA?AV?$vector@PEAVHloInstruction@xla@@V?$allocator@PEAVHloInstruction@xla@@@std@@@std@@XZ
+ ?CommittedChunksDirectory@data@tensorflow@@YA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$basic_string_view@DU?$char_traits@D@std@@@4@@Z
+ ?Compare@DefaultFieldComparator@util@protobuf@google@@UEAA?AW4ComparisonResult@FieldComparator@234@AEBVMessage@34@0PEBVFieldDescriptor@34@HHPEBVFieldContext@234@@Z
+ ?Compare@MessageDifferencer@util@protobuf@google@@QEAA_NAEBVMessage@34@0@Z
+ ?Compute@AsyncOpKernel@tensorflow@@UEAAXPEAVOpKernelContext@2@@Z
+ ?ComputeGradient@Tape@gradients@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@PEAVAbstractContext@3@V?$Span@QEAVAbstractTensorHandle@tensorflow@@@56@11V?$Span@PEAVAbstractTensorHandle@tensorflow@@@56@@Z
+ ?Convert@PythonTensorConverter@tensorflow@@QEBA?AV?$unique_ptr@U_object@@UPyDecrefDeleter@detail@tensorflow@@@std@@PEAU_object@@AEAW4DataType@2@PEA_N@Z
+ ?Convert@tflite@@YAPEAU_object@@PEAU2@00_N0PEBVPyFunctionLibrary@quantization@tensorflow@@@Z
+ ?ConvertPyObjectToAttributeType@tensorflow@@YA?AV?$unique_ptr@U_object@@UPyDecrefDeleter@detail@tensorflow@@@std@@PEAU_object@@W4AttributeType@1@@Z
+ ?ConvertPythonAPIParameters@tensorflow@@YA_NAEBVPythonAPIInfo@1@AEBVPythonTensorConverter@1@V?$Span@PEAU_object@@@lts_20260526@absl@@PEAUInferredAttributes@21@@Z
+ ?ConvertToEagerTensor@tensorflow@@YAPEAUTFE_TensorHandle@@PEAUTFE_Context@@PEAU_object@@W4DataType@1@PEBD@Z
+ ?Copy@container_internal@lts_20260526@absl@@YAXAEAVCommonFields@123@AEBUPolicyFunctions@123@AEBV4123@V?$FunctionRef@$$A6AXPEAXPEBX@Z@23@@Z
+ ?CopyFile@Env@tsl@@QEAA?AVStatus@lts_20260526@absl@@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@0@Z
+ ?CopyFrom@AttrValue@tensorflow@@QEAAXAEBV12@@Z
+ ?CopyFrom@ConfigProto@tensorflow@@QEAAXAEBV12@@Z
+ ?CopyFrom@DeviceProperties@tensorflow@@QEAAXAEBV12@@Z
+ ?CopyFrom@FunctionDef@tensorflow@@QEAAXAEBV12@@Z
+ ?CopyFrom@NodeDef@tensorflow@@QEAAXAEBV12@@Z
+ ?CopyFrom@StepStats@tensorflow@@QEAAXAEBV12@@Z
+ ?CopyPythonAPITensorLists@tensorflow@@YA_NAEBVPythonAPIInfo@1@V?$Span@PEAU_object@@@lts_20260526@absl@@@Z
+ ?CostGraphToOpPerformanceData@grappler@tensorflow@@YA?AVOpPerformanceList@2@AEBVCostGraphDef@2@AEBVGraphDef@2@@Z
+ ?Crash@Helper@internal_statusor@lts_20260526@absl@@SAXAEBVStatus@34@@Z
+ ?Create@PtrOpKernelFactory@OpKernelRegistrar@kernel_factory@tensorflow@@UEAAPEAVOpKernel@4@PEAVOpKernelConstruction@4@@Z
+ ?CreateDir@Env@tsl@@QEAA?AVStatus@lts_20260526@absl@@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@Z
+ ?CreateFingerprintDef@fingerprinting@saved_model@tensorflow@@YA?AV?$StatusOr@VFingerprintDef@tensorflow@@@lts_20260526@absl@@V?$basic_string_view@DU?$char_traits@D@std@@@std@@@Z
+ ?CreateMesh@Mesh@dtensor@tensorflow@@SA?AV123@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@AEBV?$vector@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$allocator@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@@5@AEBV?$vector@_JV?$allocator@_J@std@@@5@2121_N@Z
+ ?CreateRecordReaderOptions@RecordReaderOptions@io@tsl@@SA?AU123@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@Z
+ ?CreateRecordWriterOptions@RecordWriterOptions@io@tsl@@SA?AU123@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@Z
+ ?CreateStatusAndConditionallyLog@StatusBuilder@lts_20260526@absl@@CA?AVStatus@23@VSourceLocation@23@V?$unique_ptr@URep@StatusBuilder@lts_20260526@absl@@U?$default_delete@URep@StatusBuilder@lts_20260526@absl@@@std@@@std@@@Z
+ ?CtxFailure@OpKernelConstruction@tensorflow@@QEAAXPEBDHAEBVStatus@lts_20260526@absl@@@Z
+ ?CtxFailure@OpKernelContext@tensorflow@@QEAAXPEBDHAEBVStatus@lts_20260526@absl@@@Z
+ ?CtxFailureWithWarning@OpKernelConstruction@tensorflow@@QEAAXPEBDHAEBVStatus@lts_20260526@absl@@@Z
+ ?CtxFailureWithWarning@OpKernelContext@tensorflow@@QEAAXPEBDHAEBVStatus@lts_20260526@absl@@@Z
+ ?DataTypeSize@tensorflow@@YAHW4DataType@1@@Z
+ ?DataTypeString@tensorflow@@YA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@W4DataType@1@@Z
+ ?DataTypeToPyObject@tensorflow@@YA?AV?$unique_ptr@U_object@@UPyDecrefDeleter@detail@tensorflow@@@std@@AEBW4DataType@1@@Z
+ ?DataType_internal_data_@tensorflow@@3QBIB
+ ?DebugInfo@PythonAPIInfo@tensorflow@@QEBA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@XZ
+ ?DebugString@CheckpointReader@checkpoint@tensorflow@@QEBA?BV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@XZ
+ ?DebugString@InferenceContext@shape_inference@tensorflow@@QEAA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@VShapeHandle@23@@Z
+ ?DebugString@InferenceContext@shape_inference@tensorflow@@QEBA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@XZ
+ ?DebugString@PyListChecker@py_dispatch@tensorflow@@UEBA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@XZ
+ ?DebugString@PySignatureChecker@py_dispatch@tensorflow@@QEBA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@XZ
+ ?DebugString@PyUnionChecker@py_dispatch@tensorflow@@UEBA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@XZ
+ ?DebugString@PythonAPIDispatcher@py_dispatch@tensorflow@@QEBA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@XZ
+ ?DebugString@Tensor@tensorflow@@QEBA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@H@Z
+ ?DebugString@TensorShapeRep@tensorflow@@QEBA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@XZ
+ ?Default@Env@tsl@@SAPEAV12@XZ
+ ?DefaultProtocol@data@tensorflow@@YA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@XZ
+ ?DeleteFile@Env@tsl@@QEAA?AVStatus@lts_20260526@absl@@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@Z
+ ?DeleteProfiler@tfprof@tensorflow@@YAXXZ
+ ?DeleteRecursively@Env@tsl@@QEAA?AVStatus@lts_20260526@absl@@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@PEA_J1@Z
+ ?Destroy@StatusBuilder@lts_20260526@absl@@CAXV?$unique_ptr@URep@StatusBuilder@lts_20260526@absl@@U?$default_delete@URep@StatusBuilder@lts_20260526@absl@@@std@@@std@@@Z
+ ?DestroyCordSlow@Cord@lts_20260526@absl@@AEAAXXZ
+ ?DestroyEagerContextThreadLocalData@tensorflow@@YAXPEAU_object@@@Z
+ ?DestructorOutOfLine@TensorShapeRep@tensorflow@@AEAAXXZ
+ ?DetailedStatsEnabled@Cluster@grappler@tensorflow@@QEBA_NXZ
+ ?Device@KernelDefBuilder@tensorflow@@QEAAAEAV12@PEBD@Z
+ ?DisableDetailedStats@Cluster@grappler@tensorflow@@QEAAX_N@Z
+ ?Dispatch@PythonAPIDispatcher@py_dispatch@tensorflow@@QEAA?AV?$unique_ptr@U_object@@UPyDecrefDeleter@detail@tensorflow@@@std@@PEAU_object@@0@Z
+ ?DivNoNan@ops@tensorflow@@YA?AVStatus@lts_20260526@absl@@PEAVAbstractContext@2@QEAVAbstractTensorHandle@2@1PEAPEAV72@PEBD3@Z
+ ?DivNoNanRegisterer@gradients@tensorflow@@YAPEAVGradientFunction@12@AEBUForwardOperation@12@@Z
+ ?Divide@InferenceContext@shape_inference@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@VDimensionHandle@23@UDimensionOrConstant@23@_NPEAV723@@Z
+ ?DoQuantizeTrainingOnSerializedGraphDef@tensorflow@@YA?AVStatus@lts_20260526@absl@@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@H0PEAV56@@Z
+ ?Doc@OpDefBuilder@tensorflow@@QEAAAEAV12@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@Z
+ ?EagerTensorFromHandle@@YAPEAU_object@@PEAUTFE_TensorHandle@@_N@Z
+ ?EagerTensor_CheckExact@@YA_NPEBU_object@@@Z
+ ?EagerTensor_Handle@@YAPEAUTFE_TensorHandle@@PEBU_object@@@Z
+ ?EinsumShape@shape_inference@tensorflow@@YA?AVStatus@lts_20260526@absl@@PEAVInferenceContext@12@@Z
+ ?EnableOpDeterminism@tsl@@YAX_N@Z
+ ?EncodeToProto@ThreadSafeHistogram@histogram@tsl@@QEBAXPEAVHistogramProto@tensorflow@@_N@Z
+ ?EnqueueOpsFanin@GrapplerItem@grappler@tensorflow@@QEBA?AV?$vector@PEBVNodeDef@tensorflow@@V?$allocator@PEBVNodeDef@tensorflow@@@std@@@std@@XZ
+ ?EnsureInitialized@DataServiceDispatcherClient@data@tensorflow@@MEAA?AVStatus@lts_20260526@absl@@XZ
+ ?Enter@PyContextManager@tensorflow@@QEAA_NPEAU_object@@@Z
+ ?EqualAttrValueWrapper@tensorflow@@YA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@AEBV23@0@Z
+ ?EqualGraphDefWrapper@tensorflow@@YA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@AEBV23@0@Z
+ ?EraseMetaOnlyLarge@container_internal@lts_20260526@absl@@YAXAEAVCommonFields@123@PEBW4ctrl_t@123@_K@Z
+ ?EraseMetaOnlySmall@container_internal@lts_20260526@absl@@YAXAEAVCommonFields@123@_N_K@Z
+ ?ExpRegisterer@gradients@tensorflow@@YAPEAVGradientFunction@12@AEBUForwardOperation@12@@Z
+ ?ExperimentalClearDefaultLayout@dtensor@tensorflow@@YAXPEAXPEAUTSL_Status@@@Z
+ ?ExperimentalClearDefaultMesh@dtensor@tensorflow@@YAXPEAXPEAUTSL_Status@@@Z
+ ?ExperimentalConvertSavedModelToMlir@tensorflow@@YA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@AEBV23@0_NPEAUTSL_Status@@@Z
+ ?ExperimentalConvertSavedModelV1ToMlir@tensorflow@@YA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@AEBV23@00_N111PEAUTSL_Status@@@Z
+ ?ExperimentalConvertSavedModelV1ToMlirLite@tensorflow@@YA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@AEBV23@00_N1PEAUTSL_Status@@@Z
+ ?ExperimentalRunPassPipeline@tensorflow@@YA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@AEBV23@0_NPEAUTSL_Status@@@Z
+ ?ExperimentalSetDefaultLayout@dtensor@tensorflow@@YAXAEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@PEAXPEAUTSL_Status@@@Z
+ ?ExperimentalSetDefaultMesh@dtensor@tensorflow@@YAXAEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@PEAXPEAUTSL_Status@@@Z
+ ?ExperimentalWriteBytecode@tensorflow@@YAXAEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@0PEAUTSL_Status@@@Z
+ ?ExportToTensorBoard@ProfilerSessionWrapper@pywrap@profiler@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@XZ
+ ?ExtendSession@tensorflow@@YAXPEAUTF_Session@@PEAUTSL_Status@@@Z
+ ?FastIntToBuffer@numbers_internal@lts_20260526@absl@@YAPEADHPEAD@Z
+ ?FastIntToBuffer@numbers_internal@lts_20260526@absl@@YAPEAD_JPEAD@Z
+ ?FastIntToBuffer@numbers_internal@lts_20260526@absl@@YAPEAD_KPEAD@Z
+ ?FetchLayout@dtensor@tensorflow@@YA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@PEAUTFE_Context@@PEAUTFE_TensorHandle@@PEAXPEAUTSL_Status@@@Z
+ ?FileExists@Env@tsl@@QEAA?AVStatus@lts_20260526@absl@@V?$basic_string_view@DU?$char_traits@D@std@@@std@@@Z
+ ?FileName@EventsWriter@tensorflow@@QEAA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@XZ
+ ?Find@ByChar@lts_20260526@absl@@QEBA?AV?$basic_string_view@DU?$char_traits@D@std@@@std@@V45@_K@Z
+ ?FindKernelDef@tensorflow@@YA?AVStatus@lts_20260526@absl@@AEBVDeviceType@tsl@@AEBVNodeDef@1@PEAPEBVKernelDef@1@PEAV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@Z
+ ?FlatBufferFileToMlir@tflite@@YA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@AEBV23@_N@Z
+ ?Flatten@swig@tensorflow@@YAPEAU_object@@PEAU3@_N@Z
+ ?FlattenDictItems@tensorflow@@YAPEAU_object@@PEAU2@@Z
+ ?FlattenForData@swig@tensorflow@@YAPEAU_object@@PEAU3@@Z
+ ?Flush@EventsWriter@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@XZ
+ ?Flush@LogMessage@log_internal@lts_20260526@absl@@QEAAXXZ
+ ?Flush@RecordWriter@io@tsl@@QEAA?AVStatus@lts_20260526@absl@@XZ
+ ?FlushExecutionFiles@DebugEventsWriter@tfdbg@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@XZ
+ ?FlushNonExecutionFiles@DebugEventsWriter@tfdbg@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@XZ
+ ?ForEachPayload@StatusRep@status_internal@lts_20260526@absl@@QEBAXV?$FunctionRef@$$A6AXV?$basic_string_view@DU?$char_traits@D@std@@@std@@AEBVCord@lts_20260526@absl@@@Z@34@@Z
+ ?FormatConvertImpl@str_format_internal@lts_20260526@absl@@YA?AU?$ArgConvertResult@$0IAAAE@@123@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@VFormatConversionSpecImpl@123@PEAVFormatSinkImpl@123@@Z
+ ?FormatPack@str_format_internal@lts_20260526@absl@@YA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@VUntypedFormatSpecImpl@123@V?$Span@$$CBVFormatArgImpl@str_format_internal@lts_20260526@absl@@@23@@Z
+ ?ForwardNonIntrinsicAttributes@util@tfg@mlir@@YAXPEAVOperation@3@0@Z
+ ?FromProto@CompileOptions@xla@@SA?AV?$StatusOr@UCompileOptions@xla@@@lts_20260526@absl@@AEBVCompileOptionsProto@2@@Z
+ ?FromProto@Layout@dtensor@tensorflow@@SA?AV?$StatusOr@VLayout@dtensor@tensorflow@@@lts_20260526@absl@@AEBVLayoutProto@23@@Z
+ ?FromString@Layout@dtensor@tensorflow@@SA?AV?$StatusOr@VLayout@dtensor@tensorflow@@@lts_20260526@absl@@V?$basic_string_view@DU?$char_traits@D@std@@@std@@@Z
+ ?FromString@Mesh@dtensor@tensorflow@@SA?AV?$StatusOr@VMesh@dtensor@tensorflow@@@lts_20260526@absl@@V?$basic_string_view@DU?$char_traits@D@std@@@std@@@Z
+ ?GenerateReport@CostAnalyzer@grappler@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@AEAV?$basic_ostream@DU?$char_traits@D@std@@@std@@_N1@Z
+ ?GenerateReport@ModelAnalyzer@grappler@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@_N0AEAV?$basic_ostream@DU?$char_traits@D@std@@@std@@@Z
+ ?GenericProtoCast@pybind11_protobuf@@YA?AVhandle@pybind11@@PEAVMessage@protobuf@google@@W4return_value_policy@3@V23@_N@Z
+ ?Get@TFE_TensorHandleCache@tensorflow@@SAPEAV12@XZ
+ ?GetAnyDeviceDetails@DeviceFactory@tensorflow@@SA?AVStatus@lts_20260526@absl@@HPEAV?$unordered_map@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V12@U?$hash@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@U?$equal_to@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@V?$allocator@U?$pair@$$CBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V12@@std@@@2@@std@@@Z
+ ?GetChildren@Env@tsl@@QEAA?AVStatus@lts_20260526@absl@@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@PEAV?$vector@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$allocator@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@@7@@Z
+ ?GetCompilerIr@tensorflow@@YA?AV?$StatusOr@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@lts_20260526@absl@@W4IrExportStage@1@PEAVProcessFunctionLibraryRuntime@1@V?$basic_string_view@DU?$char_traits@D@std@@@std@@2PEAVEagerContext@1@V?$Span@$$CBUArgShapeAndDType@tensorflow@@@34@V?$Span@QEBVTensorHandle@tensorflow@@@34@W4CompilerArgSource@1@@Z
+ ?GetCompilerIr@tensorflow@@YA?AV?$StatusOr@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@lts_20260526@absl@@W4IrExportStage@1@PEAVProcessFunctionLibraryRuntime@1@V?$basic_string_view@DU?$char_traits@D@std@@@std@@PEAVDevice@1@PEAVEagerContext@1@V?$Span@$$CBUArgShapeAndDType@tensorflow@@@34@V?$Span@QEBVTensorHandle@tensorflow@@@34@W4CompilerArgSource@1@@Z
+ ?GetCurrentTimeNanos@profiler@tsl@@YA_JXZ
+ ?GetDataServiceMetadata@DataServiceDispatcherClient@data@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@AEAVDataServiceMetadata@23@@Z
+ ?GetDebugEventsWriter@DebugEventsWriter@tfdbg@tensorflow@@SAPEAV123@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@0_J@Z
+ ?GetDeviceInfo@OpLevelCostEstimator@grappler@tensorflow@@UEBA?AUDeviceInfo@23@AEBVDeviceProperties@3@@Z
+ ?GetDeviceInfo@grappler@tensorflow@@YA?AVDeviceProperties@2@AEBUParsedName@DeviceNameUtils@tsl@@@Z
+ ?GetEagerContextThreadLocalData@tensorflow@@YAPEAUEagerContextThreadLocalData@1@PEAU_object@@@Z
+ ?GetFileSize@Env@tsl@@QEAA?AVStatus@lts_20260526@absl@@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@PEA_K@Z
+ ?GetHandleShapeAndType@tensorflow@@YA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@PEAUTF_Graph@@UTF_Output@@@Z
+ ?GetLayout@Layout@dtensor@tensorflow@@SA?AV?$StatusOr@VLayout@dtensor@tensorflow@@@lts_20260526@absl@@W4LayoutType@123@AEBV?$vector@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$allocator@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@@std@@AEBVMesh@23@@Z
+ ?GetLayouts@SparseCoreLayoutStacker@tpu@tensorflow@@QEAA?AV?$StatusOr@VSparseCoreTableLayouts@tpu@tensorflow@@@lts_20260526@absl@@XZ
+ ?GetLinkedTensorRTVersion@tensorrt@tensorflow@@YA?AV?$tuple@HHH@std@@XZ
+ ?GetLoadedTensorRTVersion@tensorrt@tensorflow@@YA?AV?$tuple@HHH@std@@XZ
+ ?GetMetadata@FunctionDef@tensorflow@@QEBA?AUMetadata@protobuf@google@@XZ
+ ?GetMetadata@FunctionDefLibrary@tensorflow@@QEBA?AUMetadata@protobuf@google@@XZ
+ ?GetMetadata@GraphDef@tensorflow@@QEBA?AUMetadata@protobuf@google@@XZ
+ ?GetMetadata@LayoutProto@dtensor@tensorflow@@QEBA?AUMetadata@protobuf@google@@XZ
+ ?GetMetadata@MeshProto@dtensor@tensorflow@@QEBA?AUMetadata@protobuf@google@@XZ
+ ?GetMetadata@MetaGraphDef@tensorflow@@QEBA?AUMetadata@protobuf@google@@XZ
+ ?GetMlirCommonFlags@tensorflow@@YAPEAUMlirCommonFlags@1@XZ
+ ?GetNodeAttr@tensorflow@@YA?AVStatus@lts_20260526@absl@@AEBVAttrSlice@1@V?$basic_string_view@DU?$char_traits@D@std@@@std@@PEAH@Z
+ ?GetNodeAttr@tensorflow@@YA?AVStatus@lts_20260526@absl@@AEBVAttrSlice@1@V?$basic_string_view@DU?$char_traits@D@std@@@std@@PEAM@Z
+ ?GetNodeAttr@tensorflow@@YA?AVStatus@lts_20260526@absl@@AEBVAttrSlice@1@V?$basic_string_view@DU?$char_traits@D@std@@@std@@PEAV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@7@@Z
+ ?GetNodeAttr@tensorflow@@YA?AVStatus@lts_20260526@absl@@AEBVAttrSlice@1@V?$basic_string_view@DU?$char_traits@D@std@@@std@@PEAV?$vector@VPartialTensorShape@tensorflow@@V?$allocator@VPartialTensorShape@tensorflow@@@std@@@7@@Z
+ ?GetNodeAttr@tensorflow@@YA?AVStatus@lts_20260526@absl@@AEBVAttrSlice@1@V?$basic_string_view@DU?$char_traits@D@std@@@std@@PEAV?$vector@VTensorShape@tensorflow@@V?$allocator@VTensorShape@tensorflow@@@std@@@7@@Z
+ ?GetNodeAttr@tensorflow@@YA?AVStatus@lts_20260526@absl@@AEBVAttrSlice@1@V?$basic_string_view@DU?$char_traits@D@std@@@std@@PEAV?$vector@_JV?$allocator@_J@std@@@7@@Z
+ ?GetNodeAttr@tensorflow@@YA?AVStatus@lts_20260526@absl@@AEBVAttrSlice@1@V?$basic_string_view@DU?$char_traits@D@std@@@std@@PEAVTensorShape@1@@Z
+ ?GetNodeAttr@tensorflow@@YA?AVStatus@lts_20260526@absl@@AEBVAttrSlice@1@V?$basic_string_view@DU?$char_traits@D@std@@@std@@PEAW4DataType@1@@Z
+ ?GetNodeAttr@tensorflow@@YA?AVStatus@lts_20260526@absl@@AEBVAttrSlice@1@V?$basic_string_view@DU?$char_traits@D@std@@@std@@PEA_J@Z
+ ?GetNodeAttr@tensorflow@@YA?AVStatus@lts_20260526@absl@@AEBVAttrSlice@1@V?$basic_string_view@DU?$char_traits@D@std@@@std@@PEA_N@Z
+ ?GetNumAvailableGPUs@grappler@tensorflow@@YAHAEBU?$pair@HH@std@@@Z
+ ?GetNumAvailableLogicalCPUCores@grappler@tensorflow@@YAHXZ
+ ?GetOperationInputs@tensorflow@@YA?AV?$vector@UTF_Output@@V?$allocator@UTF_Output@@@std@@@std@@PEAUTF_Operation@@@Z
+ ?GetOutputDimensions@PjRtExecutable@xla@@UEBA?AV?$StatusOr@V?$vector@V?$vector@V?$InlinedVector@_J$05V?$allocator@_J@std@@@lts_20260526@absl@@V?$allocator@V?$InlinedVector@_J$05V?$allocator@_J@std@@@lts_20260526@absl@@@std@@@std@@V?$allocator@V?$vector@V?$InlinedVector@_J$05V?$allocator@_J@std@@@lts_20260526@absl@@V?$allocator@V?$InlinedVector@_J$05V?$allocator@_J@std@@@lts_20260526@absl@@@std@@@std@@@2@@std@@@lts_20260526@absl@@XZ
+ ?GetOutputElementTypes@PjRtExecutable@xla@@UEBA?AV?$StatusOr@V?$vector@V?$vector@W4PrimitiveType@xla@@V?$allocator@W4PrimitiveType@xla@@@std@@@std@@V?$allocator@V?$vector@W4PrimitiveType@xla@@V?$allocator@W4PrimitiveType@xla@@@std@@@std@@@2@@std@@@lts_20260526@absl@@XZ
+ ?GetOutputLayouts@PjRtExecutable@xla@@UEBA?AV?$StatusOr@V?$vector@V?$shared_ptr@$$CBVPjRtLayout@xla@@@std@@V?$allocator@V?$shared_ptr@$$CBVPjRtLayout@xla@@@std@@@2@@std@@@lts_20260526@absl@@XZ
+ ?GetOutputProperties@GraphProperties@grappler@tensorflow@@QEBAAEBV?$vector@VOpInfo_TensorProperties@tensorflow@@V?$allocator@VOpInfo_TensorProperties@tensorflow@@@std@@@std@@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@5@@Z
+ ?GetOutputShapes@PjRtExecutable@xla@@UEBA?AV?$StatusOr@V?$vector@VShape@xla@@V?$allocator@VShape@xla@@@std@@@std@@@lts_20260526@absl@@XZ
+ ?GetOutputShardings@PjRtExecutable@xla@@UEBA?AV?$optional@V?$vector@VOpSharding@xla@@V?$allocator@VOpSharding@xla@@@std@@@std@@@std@@XZ
+ ?GetOutputString@StatsCalculator@tsl@@QEBA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@XZ
+ ?GetParameterLayouts@PjRtExecutable@xla@@UEBA?AV?$StatusOr@V?$vector@V?$shared_ptr@$$CBVPjRtLayout@xla@@@std@@V?$allocator@V?$shared_ptr@$$CBVPjRtLayout@xla@@@std@@@2@@std@@@lts_20260526@absl@@XZ
+ ?GetParameterShardings@PjRtExecutable@xla@@UEBA?AV?$optional@V?$vector@VOpSharding@xla@@V?$allocator@VOpSharding@xla@@@std@@@std@@@std@@XZ
+ ?GetPythonWrappers@tensorflow@@YA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@PEBD_K@Z
+ ?GetRefForEmptyClass@container_internal@lts_20260526@absl@@YAPEAXAEAVCommonFields@123@@Z
+ ?GetRegisteredOpConverters@tensorrt@tensorflow@@YA?AV?$vector@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$allocator@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@@std@@XZ
+ ?GetRegisteredOps@OpRegistry@tensorflow@@QEAAXPEAV?$vector@VOpDef@tensorflow@@V?$allocator@VOpDef@tensorflow@@@std@@@std@@@Z
+ ?GetRegisteredXlaOpsForDevice@tensorflow@@YA?AV?$StatusOr@V?$vector@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$allocator@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@@std@@@lts_20260526@absl@@V?$basic_string_view@DU?$char_traits@D@std@@@std@@@Z
+ ?GetScalarFromTensor@InferenceContext@shape_inference@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@PEBVTensor@3@_JPEA_J@Z
+ ?GetSingleDeviceMesh@Mesh@dtensor@tensorflow@@SA?AV?$StatusOr@VMesh@dtensor@tensorflow@@@lts_20260526@absl@@V?$basic_string_view@DU?$char_traits@D@std@@@std@@@Z
+ ?GetSnapshotSplit@DataServiceDispatcherClient@data@tensorflow@@UEAA?AVStatus@lts_20260526@absl@@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@0_J11AEAVTensor@3@AEA_JAEA_N@Z
+ ?GetStats@dtensor@tensorflow@@YA?AV?$unordered_map@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@HU?$hash@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@U?$equal_to@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@V?$allocator@U?$pair@$$CBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@H@std@@@2@@std@@PEAUTFE_Context@@PEAXPEAUTSL_Status@@@Z
+ ?GetStatus@Impl@MakeErrorStream@status_macros@xla@@QEAA?AVStatus@lts_20260526@absl@@XZ
+ ?GetTableStacks@tensorflow@@YA?AV?$vector@V?$vector@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$allocator@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@@std@@V?$allocator@V?$vector@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$allocator@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@@std@@@2@@std@@AEBV?$vector@_JV?$allocator@_J@std@@@3@000AEBV?$vector@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$allocator@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@@3@_J@Z
+ ?GetTensor@CheckpointReader@checkpoint@tensorflow@@QEBAXAEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@PEAV?$unique_ptr@VTensor@tensorflow@@U?$default_delete@VTensor@tensorflow@@@std@@@5@PEAUTSL_Status@@@Z
+ ?GetUserFrames@FrozenStackTrace@tensorflow@@UEBA?AV?$vector@UStackFrame@tsl@@V?$allocator@UStackFrame@tsl@@@std@@@std@@H@Z
+ ?GetVariableToDataTypeMap@CheckpointReader@checkpoint@tensorflow@@QEBAAEBV?$unordered_map@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@W4DataType@tensorflow@@U?$hash@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@U?$equal_to@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@V?$allocator@U?$pair@$$CBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@W4DataType@tensorflow@@@std@@@2@@std@@XZ
+ ?GetVariableToShapeMap@CheckpointReader@checkpoint@tensorflow@@QEBAAEBV?$unordered_map@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@VTensorShape@tensorflow@@U?$hash@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@U?$equal_to@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@V?$allocator@U?$pair@$$CBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@VTensorShape@tensorflow@@@std@@@2@@std@@XZ
+ ?GetXlaDeviceFlags@tensorflow@@YAPEAUXlaDeviceFlags@1@XZ
+ ?Global@OpRegistry@tensorflow@@SAPEAV12@XZ
+ ?Global@flags@tensorflow@@YAAEAVFlags@12@XZ
+ ?GlobalShapeFromLocalShape@Layout@dtensor@tensorflow@@QEBA?AV?$vector@_JV?$allocator@_J@std@@@std@@V?$Span@$$CB_J@lts_20260526@absl@@PEBV?$vector@V?$vector@_JV?$allocator@_J@std@@@std@@V?$allocator@V?$vector@_JV?$allocator@_J@std@@@std@@@2@@5@@Z
+ ?GpuSupportsHalfMatMulAndConv@tensorflow@@YA_NXZ
+ ?GrapplerItemFromMetaGraphDef@grappler@tensorflow@@YA?AV?$unique_ptr@UGrapplerItem@grappler@tensorflow@@U?$default_delete@UGrapplerItem@grappler@tensorflow@@@std@@@std@@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@4@AEBVMetaGraphDef@2@AEBUItemConfig@12@@Z
+ ?HandleInvalidStatusCtorArg@Helper@internal_statusor@lts_20260526@absl@@SAXPEAVStatus@34@@Z
+ ?HasAtomicMove@Env@tsl@@QEAA?AVStatus@lts_20260526@absl@@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@PEA_N@Z
+ ?HasTensor@CheckpointReader@checkpoint@tensorflow@@QEBA_NAEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@Z
+ ?Hash64@tsl@@YA_KPEBD_K1@Z
+ ?ImportFunction@tensorflow@@YA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@AEBV23@0_NPEAUTFE_Context@@PEAUTSL_Status@@@Z
+ ?ImportGraphDef@tensorflow@@YA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@AEBV23@0_NPEAUTSL_Status@@@Z
+ ?ImportGraphDef@tensorflow@@YA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@AEBV23@0_NV?$basic_string_view@DU?$char_traits@D@std@@@3@222PEAUTSL_Status@@@Z
+ ?ImportNumpy@tsl@@YAXXZ
+ ?ImportStatusModule@google@pybind11@@YA?AVmodule_@2@_N@Z
+ ?InferDynamically@GraphMemory@grappler@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@PEAVCluster@23@@Z
+ ?InferStatically@GraphMemory@grappler@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@AEBV?$unordered_map@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@VDeviceProperties@tensorflow@@U?$hash@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@U?$equal_to@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@V?$allocator@U?$pair@$$CBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@VDeviceProperties@tensorflow@@@std@@@2@@std@@@Z
+ ?InferStatically@GraphProperties@grappler@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@_N000@Z
+ ?InfoAboutUnusedCPUFeatures@port@tensorflow@@YAXXZ
+ ?Init@DebugEventsWriter@tfdbg@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@XZ
+ ?Init@PyExceptionRegistry@tensorflow@@SAXPEAU_object@@@Z
+ ?InitInternal@OpKernelRegistrar@kernel_factory@tensorflow@@AEAAXPEBVKernelDef@3@V?$basic_string_view@DU?$char_traits@D@std@@@std@@V?$unique_ptr@VOpKernelFactory@kernel_factory@tensorflow@@U?$default_delete@VOpKernelFactory@kernel_factory@tensorflow@@@std@@@6@@Z
+ ?InitRepImpl@StatusBuilder@lts_20260526@absl@@CAPEAURep@123@VStatus@23@@Z
+ ?InitWithSuffix@EventsWriter@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@Z
+ ?Initialize@AnalyticalCostEstimator@grappler@tensorflow@@UEAA?AVStatus@lts_20260526@absl@@AEBUGrapplerItem@23@@Z
+ ?Initialize@DataServiceDispatcherClient@data@tensorflow@@UEAA?AVStatus@lts_20260526@absl@@XZ
+ ?Initialize@MeasuringCostEstimator@grappler@tensorflow@@UEAA?AVStatus@lts_20260526@absl@@AEBUGrapplerItem@23@@Z
+ ?InitializeFromParamSpecs@PythonAPIInfo@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@AEBV?$map@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V12@U?$less@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@V?$allocator@U?$pair@$$CBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V12@@std@@@2@@std@@0V?$vector@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$allocator@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@@7@PEAU_object@@@Z
+ ?InitializeFromRegisteredOp@PythonAPIInfo@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@Z
+ ?InitializePyTrampoline@tensorflow@@YAXPEAU_object@@@Z
+ ?InitializePybindProtoCastUtil@pybind11_protobuf@@YAXXZ
+ ?Input@OpDefBuilder@tensorflow@@QEAAAEAV12@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@Z
+ ?InstallStacktraceHandler@testing@tsl@@YAXXZ
+ ?IsAArch32Available@tensorflow@@YA_NXZ
+ ?IsAArch64Available@tensorflow@@YA_NXZ
+ ?IsAlreadyExists@lts_20260526@absl@@YA_NAEBVStatus@12@@Z
+ ?IsAttrs@swig@tensorflow@@YA_NPEAU_object@@@Z
+ ?IsBatchParallel@Layout@dtensor@tensorflow@@QEBA_NXZ
+ ?IsBuiltWithNvcc@tensorflow@@YA_NXZ
+ ?IsBuiltWithROCm@tensorflow@@YA_NXZ
+ ?IsBuiltWithXLA@tensorflow@@YA_NXZ
+ ?IsCompositeTensor@swig@tensorflow@@YA_NPEAU_object@@@Z
+ ?IsCoreTypeValue@swig@tensorflow@@YA_NPEAU_object@@@Z
+ ?IsDTensor@dtensor@tensorflow@@YA_NPEAUTFE_Context@@PEAUTFE_TensorHandle@@PEAXPEAUTSL_Status@@@Z
+ ?IsDataTypeSupportedByOneDNNOnThisCPU@tensorflow@@YA_NAEBW4DataType@1@@Z
+ ?IsDirectory@Env@tsl@@QEAA?AVStatus@lts_20260526@absl@@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@Z
+ ?IsEagerTensorSlow@swig@tensorflow@@YA_NPEAU_object@@@Z
+ ?IsFailedPrecondition@lts_20260526@absl@@YA_NAEBVStatus@12@@Z
+ ?IsFullyReplicated@Layout@dtensor@tensorflow@@QEBA_NXZ
+ ?IsGoogleCudaEnabled@tensorflow@@YA_NXZ
+ ?IsGoogleTensorRTEnabled@tensorrt@tensorflow@@YA_NXZ
+ ?IsMapping@swig@tensorflow@@YA_NPEAU_object@@@Z
+ ?IsMappingView@swig@tensorflow@@YA_NPEAU_object@@@Z
+ ?IsMeshDim@Mesh@dtensor@tensorflow@@QEBA_NAEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@Z
+ ?IsMklEnabled@tensorflow@@YA_NXZ
+ ?IsMutableMapping@swig@tensorflow@@YA_NPEAU_object@@@Z
+ ?IsNamedtuple@swig@tensorflow@@YAPEAU_object@@PEAU3@_N@Z
+ ?IsNested@swig@tensorflow@@YA_NPEAU_object@@@Z
+ ?IsNestedForData@swig@tensorflow@@YA_NPEAU_object@@@Z
+ ?IsNestedOrComposite@swig@tensorflow@@YA_NPEAU_object@@@Z
+ ?IsOutOfRange@lts_20260526@absl@@YA_NAEBVStatus@12@@Z
+ ?IsPowerPCAvailable@tensorflow@@YA_NXZ
+ ?IsResourceVariable@swig@tensorflow@@YA_NPEAU_object@@@Z
+ ?IsSameAddressSpace@DeviceNameUtils@tsl@@SA_NV?$basic_string_view@DU?$char_traits@D@std@@@std@@0@Z
+ ?IsSameSize@TensorShape@tensorflow@@QEBA_NAEBV12@@Z
+ ?IsSparseDTensor@dtensor@tensorflow@@YA_NPEAUTFE_Context@@PEAUTFE_TensorHandle@@PEAXPEAUTSL_Status@@@Z
+ ?IsSystemZAvailable@tensorflow@@YA_NXZ
+ ?IsTensor@swig@tensorflow@@YA_NPEAU_object@@@Z
+ ?IsTensorProtocol@swig@tensorflow@@YA_NPEAU_object@@@Z
+ ?IsTensorSpec@swig@tensorflow@@YA_NPEAU_object@@@Z
+ ?IsTypeSpec@swig@tensorflow@@YA_NPEAU_object@@@Z
+ ?IsVariable@swig@tensorflow@@YA_NPEAU_object@@@Z
+ ?IsX86Available@tensorflow@@YA_NXZ
+ ?IterateOverFullSlots@container_internal@lts_20260526@absl@@YAXAEBVCommonFields@123@_KV?$FunctionRef@$$A6AXPEBW4ctrl_t@container_internal@lts_20260526@absl@@PEAX@Z@23@@Z
+ ?Join@GrpcDataServerBase@data@tensorflow@@QEAAXXZ
+ ?LastUserFrame@FrozenStackTrace@tensorflow@@UEBA?AUStackFrame@tsl@@XZ
+ ?ListAllPhysicalDevices@DeviceFactory@tensorflow@@SA?AVStatus@lts_20260526@absl@@PEAV?$vector@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$allocator@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@@std@@@Z
+ ?ListPluggablePhysicalDevices@DeviceFactory@tensorflow@@SA?AVStatus@lts_20260526@absl@@PEAV?$vector@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$allocator@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@@std@@@Z
+ ?LoadTracesFromDebugInfoStr@tensorflow@@YA?AV?$StatusOr@V?$flat_hash_map@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$shared_ptr@VAbstractStackTrace@tensorflow@@@2@UStringHash@container_internal@lts_20260526@absl@@UStringEq@567@V?$allocator@U?$pair@$$CBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$shared_ptr@VAbstractStackTrace@tensorflow@@@2@@std@@@2@@lts_20260526@absl@@@lts_20260526@absl@@V?$basic_string_view@DU?$char_traits@D@std@@@std@@@Z
+ ?LocalShapeFromGlobalShape@Layout@dtensor@tensorflow@@QEBA?AV?$vector@_JV?$allocator@_J@std@@@std@@V?$Span@$$CB_J@lts_20260526@absl@@@Z
+ ?Log1p@ops@tensorflow@@YA?AVStatus@lts_20260526@absl@@PEAVAbstractContext@2@QEAVAbstractTensorHandle@2@PEAPEAV72@PEBD3@Z
+ ?Log1pRegisterer@gradients@tensorflow@@YAPEAVGradientFunction@12@AEBUForwardOperation@12@@Z
+ ?LookUpDebugEventsWriter@DebugEventsWriter@tfdbg@tensorflow@@SA?AVStatus@lts_20260526@absl@@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@PEAPEAV123@@Z
+ ?LookUpOpDef@OpRegistryInterface@tensorflow@@QEBA?AVStatus@lts_20260526@absl@@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@PEAPEBVOpDef@2@@Z
+ ?Lookup@PyExceptionRegistry@tensorflow@@SAPEAU_object@@W4TSL_Code@@@Z
+ ?LoopRegionAddArgument@util@tfg@mlir@@YA?AULoopRegionArgumentUpdate@123@AEAVRegion@3@VType@3@@Z
+ ?LoopRegionResultAdded@util@tfg@mlir@@YAXAEAVRegion@3@I@Z
+ ?Main@tfcompile@tensorflow@@YA?AVStatus@lts_20260526@absl@@AEBUMainFlags@12@@Z
+ ?MainOpsFanin@GrapplerItem@grappler@tensorflow@@QEBA?AV?$vector@PEBVNodeDef@tensorflow@@V?$allocator@PEBVNodeDef@tensorflow@@@std@@@std@@XZ
+ ?MakeCheckFailString@status_internal@lts_20260526@absl@@YAPEBDPEBVStatus@23@PEBD@Z
+ ?MakeEagerContextThreadLocalData@tensorflow@@YAXPEAU_object@@00@Z
+ ?MakeFingerprintJson@metrics@tensorflow@@YA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@VFingerprintDef@2@@Z
+ ?MakeRefCountingHandle@ResourceHandle@tensorflow@@SA?AV12@PEAVResourceBase@2@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@AEBVTypeIndex@2@AEBV?$vector@UDtypeAndPartialTensorShape@tensorflow@@V?$allocator@UDtypeAndPartialTensorShape@tensorflow@@@std@@@5@AEBV?$optional@VManagedStackTrace@tensorflow@@@5@@Z
+ ?MakeRep@Status@lts_20260526@absl@@CA_K_KV?$basic_string_view@DU?$char_traits@D@std@@@std@@VSourceLocation@23@@Z
+ ?MakeSavedModelPathAndSingleprint@metrics@tensorflow@@YA?AV?$StatusOr@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@lts_20260526@absl@@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@0@Z
+ ?MakeShape@InferenceContext@shape_inference@tensorflow@@QEAA?AVShapeHandle@23@AEBV?$vector@VDimensionHandle@shape_inference@tensorflow@@V?$allocator@VDimensionHandle@shape_inference@tensorflow@@@std@@@std@@@Z
+ ?MakeShape@InferenceContext@shape_inference@tensorflow@@QEAA?AVShapeHandle@23@V?$initializer_list@UDimensionOrConstant@shape_inference@tensorflow@@@std@@@Z
+ ?MakeShapeFromPartialTensorShape@InferenceContext@shape_inference@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@AEBVPartialTensorShape@3@PEAVShapeHandle@23@@Z
+ ?MakeShapeFromShapeTensor@InferenceContext@shape_inference@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@HPEAVShapeHandle@23@@Z
+ ?MakeShapeFromTensorShape@InferenceContext@shape_inference@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@AEBVTensorShape@3@PEAVShapeHandle@23@@Z
+ ?MapToLocalCode@status_internal@lts_20260526@absl@@YA?AW4StatusCode@23@H@Z
+ ?MatMul@ops@tensorflow@@YA?AVStatus@lts_20260526@absl@@PEAVAbstractContext@2@QEAVAbstractTensorHandle@2@1PEAPEAV72@_N333PEBD4@Z
+ ?MatMulRegisterer@gradients@tensorflow@@YAPEAVGradientFunction@12@AEBUForwardOperation@12@@Z
+ ?MatchSignature@OpKernelContext@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@V?$Span@$$CBW4DataType@tensorflow@@@45@0@Z
+ ?Matrix@InferenceContext@shape_inference@tensorflow@@QEAA?AVShapeHandle@23@UDimensionOrConstant@23@0@Z
+ ?MaybeRaiseExceptionFromTFStatus@tensorflow@@YAHPEAUTSL_Status@@PEAU_object@@@Z
+ ?MemoryTypesForNode@tensorflow@@YA?AVStatus@lts_20260526@absl@@PEBVOpRegistryInterface@1@AEBVDeviceType@tsl@@AEBVNodeDef@1@PEAV?$InlinedVector@W4MemoryType@tensorflow@@$03V?$allocator@W4MemoryType@tensorflow@@@std@@@34@3@Z
+ ?Merge@InferenceContext@shape_inference@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@VDimensionHandle@23@0PEAV723@@Z
+ ?Merge@InferenceContext@shape_inference@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@VShapeHandle@23@0PEAV723@@Z
+ ?MeshDimNames@Mesh@dtensor@tensorflow@@QEBA?AV?$vector@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$allocator@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@@std@@XZ
+ ?MlirQuantizeModel@tflite@@YAPEAU_object@@PEAU2@_N1HHH1100110@Z
+ ?MlirSparsifyModel@tflite@@YAPEAU_object@@PEAU2@@Z
+ ?Mul@ops@tensorflow@@YA?AVStatus@lts_20260526@absl@@PEAVAbstractContext@2@QEAVAbstractTensorHandle@2@1PEAPEAV72@PEBD3@Z
+ ?MulRegisterer@gradients@tensorflow@@YAPEAVGradientFunction@12@AEBUForwardOperation@12@@Z
+ ?Multiply@InferenceContext@shape_inference@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@VDimensionHandle@23@UDimensionOrConstant@23@PEAV723@@Z
+ ?NameRangesForNode@tensorflow@@YA?AVStatus@lts_20260526@absl@@AEBVAttrSlice@1@AEBVOpDef@1@PEAV?$FlatMap@V?$basic_string_view@DU?$char_traits@D@std@@@std@@U?$pair@HH@2@U?$hash@V?$basic_string_view@DU?$char_traits@D@std@@@std@@X@tsl@@U?$equal_to@V?$basic_string_view@DU?$char_traits@D@std@@@std@@@2@@gtl@tsl@@2@Z
+ ?Neg@ops@tensorflow@@YA?AVStatus@lts_20260526@absl@@PEAVAbstractContext@2@QEAVAbstractTensorHandle@2@PEAPEAV72@PEBD3@Z
+ ?NegRegisterer@gradients@tensorflow@@YAPEAVGradientFunction@12@AEBUForwardOperation@12@@Z
+ ?NewAppendableFile@Env@tsl@@QEAA?AVStatus@lts_20260526@absl@@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@PEAV?$unique_ptr@VWritableFile@tsl@@U?$default_delete@VWritableFile@tsl@@@std@@@7@@Z
+ ?NewDispatchServer@data@tensorflow@@YA?AVStatus@lts_20260526@absl@@AEBVDispatcherConfig@experimental@12@AEAV?$unique_ptr@VDispatchGrpcDataServer@data@tensorflow@@U?$default_delete@VDispatchGrpcDataServer@data@tensorflow@@@std@@@std@@@Z
+ ?NewProfiler@tfprof@tensorflow@@YA_NPEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@0@Z
+ ?NewRandomAccessFile@Env@tsl@@QEAA?AVStatus@lts_20260526@absl@@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@PEAV?$unique_ptr@VRandomAccessFile@tsl@@U?$default_delete@VRandomAccessFile@tsl@@@std@@@7@@Z
+ ?NewWorkerServer@data@tensorflow@@YA?AVStatus@lts_20260526@absl@@AEBVWorkerConfig@experimental@12@AEAV?$unique_ptr@VWorkerGrpcDataServer@data@tensorflow@@U?$default_delete@VWorkerGrpcDataServer@data@tensorflow@@@std@@@std@@@Z
+ ?NewWritableFile@Env@tsl@@QEAA?AVStatus@lts_20260526@absl@@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@PEAV?$unique_ptr@VWritableFile@tsl@@U?$default_delete@VWritableFile@tsl@@@std@@@7@@Z
+ ?NumCheckpointShardsWritten@metrics@tensorflow@@YAAEAVCounterCell@monitoring@tsl@@XZ
+ ?NumTasks@WorkerGrpcDataServer@data@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@PEAH@Z
+ ?NumWorkers@DispatchGrpcDataServer@data@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@PEAH@Z
+ ?OpDeterminismRequired@tsl@@YA_NXZ
+ ?OpHasDevice@util@tfg@mlir@@YA_NPEAVOperation@3@PEBD@Z
+ ?OptimizeConsumeItem@MetaOptimizer@grappler@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@PEAVCluster@23@$$QEAUGrapplerItem@23@PEAVGraphDef@3@@Z
+ ?Output@OpDefBuilder@tensorflow@@QEAAAEAV12@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@Z
+ ?Pack@dtensor@tensorflow@@YAPEAUTFE_TensorHandle@@PEAUTFE_Context@@HPEAPEAU3@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@PEAXPEAUTSL_Status@@@Z
+ ?ParseFromProto@Mesh@dtensor@tensorflow@@SA?AV?$StatusOr@VMesh@dtensor@tensorflow@@@lts_20260526@absl@@AEBVMeshProto@23@@Z
+ ?ParseFullName@DeviceNameUtils@tsl@@SA_NV?$basic_string_view@DU?$char_traits@D@std@@@std@@PEAUParsedName@12@@Z
+ ?ParseFullOrLocalName@DeviceNameUtils@tsl@@SA_NV?$basic_string_view@DU?$char_traits@D@std@@@std@@PEAUParsedName@12@@Z
+ ?ParseSavedModelPathAndSingleprint@metrics@tensorflow@@YA?AV?$StatusOr@U?$pair@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V12@@std@@@lts_20260526@absl@@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@Z
+ ?ParseTensorName@tensorflow@@YA?AUTensorId@1@V?$basic_string_view@DU?$char_traits@D@std@@@std@@@Z
+ ?ParseTransformParameters@graph_transforms@tensorflow@@YA?AVStatus@lts_20260526@absl@@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@PEAV?$vector@U?$pair@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$map@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$vector@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$allocator@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@@2@U?$less@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@V?$allocator@U?$pair@$$CBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$vector@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$allocator@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@@2@@std@@@2@@2@@std@@V?$allocator@U?$pair@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$map@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$vector@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$allocator@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@@2@U?$less@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@V?$allocator@U?$pair@$$CBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$vector@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$allocator@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@@2@@std@@@2@@2@@std@@@2@@7@@Z
+ ?PredictCosts@AnalyticalCostEstimator@grappler@tensorflow@@UEBA?AVStatus@lts_20260526@absl@@AEBVGraphDef@3@PEAVRunMetadata@3@PEAUCosts@23@@Z
+ ?PredictCosts@MeasuringCostEstimator@grappler@tensorflow@@UEBA?AVStatus@lts_20260526@absl@@AEBVGraphDef@3@PEAVRunMetadata@3@PEAUCosts@23@@Z
+ ?PredictCosts@OpLevelCostEstimator@grappler@tensorflow@@UEBA?AUCosts@23@AEBUOpContext@23@@Z
+ ?PrepareInsertLarge@container_internal@lts_20260526@absl@@YA_KAEAVCommonFields@123@AEBUPolicyFunctions@123@_KV?$NonIterableBitMask@I$0BA@$0A@@123@UFindInfo@123@@Z
+ ?PrepareInsertSmallNonSoo@container_internal@lts_20260526@absl@@YA?AU?$pair@PEAW4ctrl_t@container_internal@lts_20260526@absl@@PEAX@std@@AEAVCommonFields@123@AEBUPolicyFunctions@123@V?$FunctionRef@$$A6A_K_K@Z@23@@Z
+ ?PrepareToModify@Status@lts_20260526@absl@@CAPEAVStatusRep@status_internal@23@_K@Z
+ ?PrintModelAnalysis@tfprof@tensorflow@@YA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@PEBV34@0000@Z
+ ?PrintResult@MetaOptimizer@grappler@tensorflow@@QEAAXXZ
+ ?PrintStepStats@StatSummarizer@tensorflow@@QEBAXXZ
+ ?ProcessStepStats@StatSummarizer@tensorflow@@QEAAXAEBVStepStats@2@@Z
+ ?Profile@tfprof@tensorflow@@YA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@PEBV34@0@Z
+ ?ProfilerFromFile@tfprof@tensorflow@@YAXPEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@Z
+ ?PyBytesAsStringView@pybind11_protobuf@@YA?AV?$basic_string_view@DU?$char_traits@D@std@@@std@@Vbytes@pybind11@@@Z
+ ?PyProtoGetCppMessagePointer@pybind11_protobuf@@YAPEBVMessage@protobuf@google@@Vhandle@pybind11@@@Z
+ ?PyProtoHasMatchingFullName@pybind11_protobuf@@YA_NVhandle@pybind11@@PEBVDescriptor@protobuf@google@@@Z
+ ?PyProtoSerializePartialToString@pybind11_protobuf@@YA?AVbytes@pybind11@@Vhandle@3@_N@Z
+ ?PywrapExpandPresets@pywrap@quantization@stablehlo@@YA?AVQuantizationConfig@23@AEBV423@@Z
+ ?PywrapPopulateDefaults@pywrap@quantization@stablehlo@@YA?AVQuantizationConfig@23@AEBV423@@Z
+ ?PywrapQuantizeStaticRangePtq@pywrap@quantization@stablehlo@@YA?AVStatus@lts_20260526@absl@@V?$basic_string_view@DU?$char_traits@D@std@@@std@@0AEBVQuantizationConfig@23@AEBV?$vector@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$allocator@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@@8@AEBV?$flat_hash_map@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@VSignatureDef@tensorflow@@UStringHash@container_internal@lts_20260526@absl@@UStringEq@678@V?$allocator@U?$pair@$$CBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@VSignatureDef@tensorflow@@@std@@@2@@56@AEBVPyFunctionLibrary@2tensorflow@@@Z
+ ?PywrapQuantizeWeightOnlyPtq@pywrap@quantization@stablehlo@@YA?AVStatus@lts_20260526@absl@@V?$basic_string_view@DU?$char_traits@D@std@@@std@@0AEBVQuantizationConfig@23@AEBV?$vector@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$allocator@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@@8@AEBV?$flat_hash_map@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@VSignatureDef@tensorflow@@UStringHash@container_internal@lts_20260526@absl@@UStringEq@678@V?$allocator@U?$pair@$$CBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@VSignatureDef@tensorflow@@@std@@@2@@56@AEBVPyFunctionLibrary@2tensorflow@@@Z
+ ?PywrapSavedModelToStablehlo@pywrap@tensorflow_to_stablehlo@mlir@@YA?AV?$StatusOr@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@lts_20260526@absl@@V?$basic_string_view@DU?$char_traits@D@std@@@std@@AEBV?$vector@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$allocator@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@@8@10@Z
+ ?PywrapTfModuleToStablehlo@pywrap@tensorflow_to_stablehlo@mlir@@YA?AV?$StatusOr@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@lts_20260526@absl@@V?$basic_string_view@DU?$char_traits@D@std@@@std@@0@Z
+ ?QuantizeDynamicRangePtq@quantization@tensorflow@@YA?AV?$StatusOr@VExportedModel@quantization@tensorflow@@@lts_20260526@absl@@V?$basic_string_view@DU?$char_traits@D@std@@@std@@AEBV?$vector@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$allocator@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@@7@AEBV?$unordered_set@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@U?$hash@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@U?$equal_to@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@V?$allocator@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@@7@AEBVQuantizationOptions@12@@Z
+ ?QuantizeQatModel@quantization@tensorflow@@YA?AV?$StatusOr@VExportedModel@quantization@tensorflow@@@lts_20260526@absl@@V?$basic_string_view@DU?$char_traits@D@std@@@std@@AEBV?$vector@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$allocator@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@@7@AEBV?$unordered_set@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@U?$hash@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@U?$equal_to@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@V?$allocator@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@@7@AEBVQuantizationOptions@12@@Z
+ ?QuantizeStaticRangePtq@quantization@tensorflow@@YA?AV?$StatusOr@VExportedModel@quantization@tensorflow@@@lts_20260526@absl@@V?$basic_string_view@DU?$char_traits@D@std@@@std@@AEBV?$vector@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$allocator@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@@7@AEBV?$unordered_set@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@U?$hash@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@U?$equal_to@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@V?$allocator@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@@7@AEBVQuantizationOptions@12@AEBV?$flat_hash_map@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@VSignatureDef@tensorflow@@UStringHash@container_internal@lts_20260526@absl@@UStringEq@678@V?$allocator@U?$pair@$$CBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@VSignatureDef@tensorflow@@@std@@@2@@45@AEBVPyFunctionLibrary@12@AEBV?$flat_hash_map@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@VRepresentativeDatasetFile@quantization@tensorflow@@UStringHash@container_internal@lts_20260526@absl@@UStringEq@789@V?$allocator@U?$pair@$$CBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@VRepresentativeDatasetFile@quantization@tensorflow@@@std@@@2@@45@@Z
+ ?QuantizeWeightOnly@quantization@tensorflow@@YA?AV?$StatusOr@VExportedModel@quantization@tensorflow@@@lts_20260526@absl@@V?$basic_string_view@DU?$char_traits@D@std@@@std@@AEBVQuantizationOptions@12@@Z
+ ?RawLog@raw_log_internal@lts_20260526@absl@@YAXW4LogSeverity@23@PEBDH1ZZ
+ ?ReadFileToString@tsl@@YA?AVStatus@lts_20260526@absl@@PEAVEnv@1@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@PEAV67@@Z
+ ?ReadLineAsString@BufferedInputStream@io@tsl@@QEAA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@XZ
+ ?ReadRecord@RecordReader@io@tsl@@QEAA?AVStatus@lts_20260526@absl@@PEA_KPEAVtstring@3@@Z
+ ?ReadSavedModelFingerprint@fingerprinting@saved_model@tensorflow@@YA?AV?$StatusOr@VFingerprintDef@tensorflow@@@lts_20260526@absl@@V?$basic_string_view@DU?$char_traits@D@std@@@std@@@Z
+ ?Record@TraceMeRecorder@profiler@tsl@@SAX$$QEAUEvent@123@@Z
+ ?RecordMutation@tensorflow@@YAXPEAUTF_Graph@@AEBUTF_Operation@@PEBD@Z
+ ?RecursivelyCreateDir@Env@tsl@@QEAA?AVStatus@lts_20260526@absl@@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@Z
+ ?ReduceScatterShape@shape_inference@tensorflow@@YA?AVStatus@lts_20260526@absl@@PEAVInferenceContext@12@@Z
+ ?Register@GradientRegistry@gradients@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$function@$$A6APEAVGradientFunction@gradients@tensorflow@@AEBUForwardOperation@23@@Z@8@@Z
+ ?Register@PythonAPIDispatcher@py_dispatch@tensorflow@@QEAAXVPySignatureChecker@23@PEAU_object@@@Z
+ ?RegisterCustomOpdefs@tflite@@YAPEAU_object@@PEAU2@@Z
+ ?RegisterDeviceAndGetId@DebugEventsWriter@tfdbg@tensorflow@@QEAAHAEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@Z
+ ?RegisterDispatchableType@py_dispatch@tensorflow@@YA_NPEAU_object@@@Z
+ ?RegisterPyObject@swig@tensorflow@@YAPEAU_object@@PEAU3@0@Z
+ ?Relu@ops@tensorflow@@YA?AVStatus@lts_20260526@absl@@PEAVAbstractContext@2@QEAVAbstractTensorHandle@2@PEAPEAV72@PEBD3@Z
+ ?ReluRegisterer@gradients@tensorflow@@YAPEAVGradientFunction@12@AEBUForwardOperation@12@@Z
+ ?RemoveControlEdge@Graph@tensorflow@@QEAAXPEBVEdge@2@@Z
+ ?RemoveNonDeprecationDescriptionsFromOpDef@tensorflow@@YAXPEAVOpDef@1@@Z
+ ?RenameFile@Env@tsl@@QEAA?AVStatus@lts_20260526@absl@@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@0@Z
+ ?ReplicatedOnMesh@Layout@dtensor@tensorflow@@SA?AV123@AEBVMesh@23@H@Z
+ ?ReserveEmptyNonAllocatedTableToFitBucketCount@container_internal@lts_20260526@absl@@YAXAEAVCommonFields@123@AEBUPolicyFunctions@123@_K@Z
+ ?ReserveTableToFitNewSize@container_internal@lts_20260526@absl@@YAXAEAVCommonFields@123@AEBUPolicyFunctions@123@_K@Z
+ ?ResetJitCompilerFlags@tensorflow@@YAXXZ
+ ?RetrieveCollectedErrors@tflite@@YA?AV?$vector@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$allocator@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@@std@@XZ
+ ?Retry@grpc_util@data@tensorflow@@YA?AVStatus@lts_20260526@absl@@AEBV?$function@$$A6A?AVStatus@lts_20260526@absl@@XZ@std@@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@8@_J@Z
+ ?RunHloCostAnalysis@PjRtExecutableUtil@xla@@SA?AV?$StatusOr@V?$flat_hash_map@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$variant@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@_N_JV?$vector@_JV?$allocator@_J@std@@@2@M@2@UStringHash@container_internal@lts_20260526@absl@@UStringEq@567@V?$allocator@U?$pair@$$CBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$variant@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@_N_JV?$vector@_JV?$allocator@_J@std@@@2@M@2@@std@@@2@@lts_20260526@absl@@@lts_20260526@absl@@AEBVPjRtExecutable@2@PEAVHloCostAnalysis@2@@Z
+ ?SameNamedtuples@swig@tensorflow@@YAPEAU_object@@PEAU3@0@Z
+ ?SavedModelFoundFingerprintOnLoad@metrics@tensorflow@@YAAEAV?$GaugeCell@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@monitoring@tsl@@XZ
+ ?SavedModelReadApi@metrics@tensorflow@@YAAEAVCounterCell@monitoring@tsl@@V?$basic_string_view@DU?$char_traits@D@std@@@std@@@Z
+ ?SavedModelReadCount@metrics@tensorflow@@YAAEAVCounterCell@monitoring@tsl@@V?$basic_string_view@DU?$char_traits@D@std@@@std@@@Z
+ ?SavedModelReadFingerprint@metrics@tensorflow@@YAAEAV?$GaugeCell@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@monitoring@tsl@@XZ
+ ?SavedModelReadPath@metrics@tensorflow@@YAAEAV?$GaugeCell@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@monitoring@tsl@@XZ
+ ?SavedModelReadPathAndSingleprint@metrics@tensorflow@@YAAEAV?$GaugeCell@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@monitoring@tsl@@XZ
+ ?SavedModelWriteApi@metrics@tensorflow@@YAAEAVCounterCell@monitoring@tsl@@V?$basic_string_view@DU?$char_traits@D@std@@@std@@@Z
+ ?SavedModelWriteCount@metrics@tensorflow@@YAAEAVCounterCell@monitoring@tsl@@V?$basic_string_view@DU?$char_traits@D@std@@@std@@@Z
+ ?SavedModelWriteFingerprint@metrics@tensorflow@@YAAEAV?$GaugeCell@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@monitoring@tsl@@XZ
+ ?SavedModelWritePath@metrics@tensorflow@@YAAEAV?$GaugeCell@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@monitoring@tsl@@XZ
+ ?SavedModelWritePathAndSingleprint@metrics@tensorflow@@YAAEAV?$GaugeCell@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@monitoring@tsl@@XZ
+ ?Scalar@InferenceContext@shape_inference@tensorflow@@QEAA?AVShapeHandle@23@XZ
+ ?Schedule@ThreadPool@thread@tsl@@QEAAXV?$function@$$A6AXXZ@std@@@Z
+ ?Seek@BufferedInputStream@io@tsl@@QEAA?AVStatus@lts_20260526@absl@@_J@Z
+ ?SerializeToString@tfprof@tensorflow@@YA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@XZ
+ ?SetAttr@tensorflow@@YAXPEAUTF_Graph@@PEAUTF_Operation@@PEBDPEAUTF_Buffer@@PEAUTSL_Status@@@Z
+ ?SetCEagerContext@tensorflow@@YAXPEAVEagerContext@1@@Z
+ ?SetDimWithStatus@?$TensorShapeBase@VPartialTensorShape@tensorflow@@@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@H_J@Z
+ ?SetFullType@tensorflow@@YAXPEAUTF_Graph@@PEAUTF_Operation@@PEBUTF_Buffer@@@Z
+ ?SetHandleShapeAndType@tensorflow@@YAXPEAUTF_Graph@@UTF_Output@@PEBX_KPEAUTSL_Status@@@Z
+ ?SetIsStateful@OpDefBuilder@tensorflow@@QEAAAEAV12@_N@Z
+ ?SetIteratorElementLayouts@dtensor@tensorflow@@YAXPEAUTFE_Context@@PEAUTFE_TensorHandle@@AEBV?$vector@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$allocator@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@@std@@PEAXPEAUTSL_Status@@@Z
+ ?SetNumWarmupSteps@Cluster@grappler@tensorflow@@QEAAXH@Z
+ ?SetPayload@StatusRep@status_internal@lts_20260526@absl@@QEAAXV?$basic_string_view@DU?$char_traits@D@std@@@std@@VCord@34@@Z
+ ?SetShapeFn@OpDefBuilder@tensorflow@@QEAAAEAV12@V?$function@$$A6A?AVStatus@lts_20260526@absl@@PEAVInferenceContext@shape_inference@tensorflow@@@Z@std@@@Z
+ ?SetTPUCoreIDs@dtensor@tensorflow@@YAXAEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@AEBV?$vector@HV?$allocator@H@std@@@4@PEAXPEAUTSL_Status@@@Z
+ ?SetVisited@?$DfsHloVisitorBase@PEAVHloInstruction@xla@@@xla@@QEAAXAEBVHloInstruction@2@@Z
+ ?SetVisited@?$DfsHloVisitorBase@PEBVHloInstruction@xla@@@xla@@QEAAXAEBVHloInstruction@2@@Z
+ ?Set_TF_Status_from_Status@tsl@@YAXPEAUTSL_Status@@AEBVStatus@lts_20260526@absl@@@Z
+ ?ShardingCallbackDescription@metrics@tensorflow@@YAAEAV?$GaugeCell@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@monitoring@tsl@@XZ
+ ?ShardingCallbackDuration@metrics@tensorflow@@YAAEAVCounterCell@monitoring@tsl@@XZ
+ ?Signal@CondVar@lts_20260526@absl@@QEAAXXZ
+ ?Singleprint@fingerprinting@saved_model@tensorflow@@YA?AV?$StatusOr@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@lts_20260526@absl@@V?$basic_string_view@DU?$char_traits@D@std@@@std@@@Z
+ ?Singleprint@fingerprinting@saved_model@tensorflow@@YA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@AEBVFingerprintDef@3@@Z
+ ?Singleprint@fingerprinting@saved_model@tensorflow@@YA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@_K000@Z
+ ?SlowCopyFrom@TensorShapeRep@tensorflow@@AEAAXAEBV12@@Z
+ ?SlowIsEnabled3@VLogSite@log_internal@lts_20260526@absl@@AEAA_NH@Z
+ ?SnapshotDoneFilePath@data@tensorflow@@YA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$basic_string_view@DU?$char_traits@D@std@@@4@@Z
+ ?SnapshotErrorFilePath@data@tensorflow@@YA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$basic_string_view@DU?$char_traits@D@std@@@4@@Z
+ ?SnapshotMetadataFilePath@data@tensorflow@@YA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$basic_string_view@DU?$char_traits@D@std@@@4@@Z
+ ?SnapshotStreams@DispatchGrpcDataServer@data@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@PEAV?$vector@USnapshotStreamInfoWrapper@data@tensorflow@@V?$allocator@USnapshotStreamInfoWrapper@data@tensorflow@@@std@@@8@@Z
+ ?SnapshotTaskProgresses@WorkerGrpcDataServer@data@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@PEAV?$vector@USnapshotTaskProgressWrapper@data@tensorflow@@V?$allocator@USnapshotTaskProgressWrapper@data@tensorflow@@@std@@@std@@@Z
+ ?SparsePack@dtensor@tensorflow@@YAPEAUTFE_TensorHandle@@PEAUTFE_Context@@HPEAPEAU3@11AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@PEAXPEAUTSL_Status@@@Z
+ ?SparseSoftmaxCrossEntropyWithLogits@ops@tensorflow@@YA?AVStatus@lts_20260526@absl@@PEAVAbstractContext@2@QEAVAbstractTensorHandle@2@1PEAPEAV72@2PEBD3@Z
+ ?SparseSoftmaxCrossEntropyWithLogitsRegisterer@gradients@tensorflow@@YAPEAVGradientFunction@12@AEBUForwardOperation@12@@Z
+ ?Start@GrpcDataServerBase@data@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@XZ
+ ?Start@ProfilerSessionWrapper@pywrap@profiler@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@PEBDAEBV?$flat_hash_map@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$variant@_NHV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@UStringHash@container_internal@lts_20260526@absl@@UStringEq@567@V?$allocator@U?$pair@$$CBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$variant@_NHV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@@std@@@2@@67@@Z
+ ?StartCancel@CancellationManager@tsl@@QEAAXXZ
+ ?StartProfilerServer@ProfilerServer@profiler@tsl@@QEAAXH@Z
+ ?Stat@Env@tsl@@QEAA?AVStatus@lts_20260526@absl@@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@PEAUFileStatistics@2@@Z
+ ?Stop@GrpcDataServerBase@data@tensorflow@@QEAAXXZ
+ ?Stop@ProfilerSessionWrapper@pywrap@profiler@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@PEAV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@Z
+ ?StrAppend@lts_20260526@absl@@YAXPEAV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@AEBVAlphaNum@12@111@Z
+ ?StrAppend@lts_20260526@absl@@YAXPEAV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@AEBVAlphaNum@12@11@Z
+ ?StrAppend@lts_20260526@absl@@YAXPEAV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@AEBVAlphaNum@12@@Z
+ ?StrCat@lts_20260526@absl@@YA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@AEBVAlphaNum@12@000@Z
+ ?StrCat@lts_20260526@absl@@YA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@AEBVAlphaNum@12@00@Z
+ ?StrCat@lts_20260526@absl@@YA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@AEBVAlphaNum@12@0@Z
+ ?StripDefaultAttributes@tensorflow@@YAXAEBVOpRegistryInterface@1@PEAV?$RepeatedPtrField@VNodeDef@tensorflow@@@protobuf@google@@@Z
+ ?Sub@ops@tensorflow@@YA?AVStatus@lts_20260526@absl@@PEAVAbstractContext@2@QEAVAbstractTensorHandle@2@1PEAPEAV72@PEBD3@Z
+ ?SubRegisterer@gradients@tensorflow@@YAPEAVGradientFunction@12@AEBUForwardOperation@12@@Z
+ ?Subtract@InferenceContext@shape_inference@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@VDimensionHandle@23@UDimensionOrConstant@23@PEAV723@@Z
+ ?TFE_DeleteContextCapsule@@YAXPEAU_object@@@Z
+ ?TFE_GetPythonString@@YAPEBDPEAU_object@@@Z
+ ?TFE_Py_EnableInteractivePythonLogging@@YAXXZ
+ ?TFE_Py_ExecuteCancelable@@YAXPEAUTFE_Context@@PEBD1PEAV?$InlinedVector@PEAUTFE_TensorHandle@@$03V?$allocator@PEAUTFE_TensorHandle@@@std@@@lts_20260526@absl@@PEAU_object@@PEAUTFE_CancellationManager@@PEAV?$InlinedVector@PEAUTFE_TensorHandle@@$01V?$allocator@PEAUTFE_TensorHandle@@@std@@@34@PEAUTSL_Status@@@Z
+ ?TFE_Py_FastPathExecute_C@@YAPEAU_object@@PEAU1@@Z
+ ?TFE_Py_ForwardAccumulatorJVP@@YAPEAU_object@@PEAU1@0@Z
+ ?TFE_Py_ForwardAccumulatorNew@@YAPEAU_object@@_N@Z
+ ?TFE_Py_ForwardAccumulatorPopState@@YAPEAU_object@@XZ
+ ?TFE_Py_ForwardAccumulatorPushState@@YAPEAU_object@@XZ
+ ?TFE_Py_ForwardAccumulatorSetAdd@@YAPEAU_object@@PEAU1@@Z
+ ?TFE_Py_ForwardAccumulatorSetRemove@@YAXPEAU_object@@@Z
+ ?TFE_Py_ForwardAccumulatorWatch@@YAXPEAU_object@@00@Z
+ ?TFE_Py_InitEagerTensor@@YAPEAU_object@@PEAU1@@Z
+ ?TFE_Py_PackJVPs@@YAPEAU_object@@PEAU1@@Z
+ ?TFE_Py_RecordGradient@@YAPEAU_object@@PEAU1@0000@Z
+ ?TFE_Py_RegisterExceptionClass@@YAPEAU_object@@PEAU1@@Z
+ ?TFE_Py_RegisterFallbackExceptionClass@@YAPEAU_object@@PEAU1@@Z
+ ?TFE_Py_RegisterGradientFunction@@YAPEAU_object@@PEAU1@@Z
+ ?TFE_Py_RegisterJVPFunction@@YAPEAU_object@@PEAU1@@Z
+ ?TFE_Py_RegisterVSpace@@YAPEAU_object@@PEAU1@@Z
+ ?TFE_Py_SetEagerContext@@YAPEAU_object@@PEAU1@@Z
+ ?TFE_Py_SetEagerTensorProfiler@@YAPEAU_object@@PEAU1@@Z
+ ?TFE_Py_TapeGradient@@YAPEAU_object@@PEAU1@00000PEAUTSL_Status@@@Z
+ ?TFE_Py_TapeSetAdd@@YAXPEAU_object@@@Z
+ ?TFE_Py_TapeSetDeleteTrace@@YAX_J@Z
+ ?TFE_Py_TapeSetIsEmpty@@YAPEAU_object@@XZ
+ ?TFE_Py_TapeSetIsStopped@@YAPEAU_object@@XZ
+ ?TFE_Py_TapeSetNew@@YAPEAU_object@@PEAU1@0@Z
+ ?TFE_Py_TapeSetPossibleGradientTypes@@YAPEAU_object@@PEAU1@@Z
+ ?TFE_Py_TapeSetRecordOperation@@YAPEAU_object@@PEAU1@0000@Z
+ ?TFE_Py_TapeSetRecordOperationBackprop@@YAPEAU_object@@PEAU1@000@Z
+ ?TFE_Py_TapeSetRecordOperationForwardprop@@YAPEAU_object@@PEAU1@0000@Z
+ ?TFE_Py_TapeSetRemove@@YAXPEAU_object@@@Z
+ ?TFE_Py_TapeSetRestartOnThread@@YAXXZ
+ ?TFE_Py_TapeSetShouldRecordBackprop@@YAPEAU_object@@PEAU1@@Z
+ ?TFE_Py_TapeSetStopOnThread@@YAXXZ
+ ?TFE_Py_TapeVariableAccessed@@YAXPEAU_object@@@Z
+ ?TFE_Py_TapeWatch@@YAXPEAU_object@@0@Z
+ ?TFE_Py_TapeWatchVariable@@YAXPEAU_object@@0@Z
+ ?TFE_Py_TapeWatchedVariables@@YAPEAU_object@@PEAU1@@Z
+ ?TFE_Py_TensorShapeOnDevice@@YAPEAU_object@@PEAU1@@Z
+ ?TFE_Py_TensorShapeSlice@@YAPEAU_object@@PEAU1@H@Z
+ ?TFE_Py_UID@@YAPEAU_object@@XZ
+ ?TFE_Py_VariableWatcherNew@@YAPEAU_object@@XZ
+ ?TFE_Py_VariableWatcherRemove@@YAXPEAU_object@@@Z
+ ?TFE_Py_VariableWatcherVariableAccessed@@YAXPEAU_object@@@Z
+ ?TFE_Py_VariableWatcherWatchedVariables@@YAPEAU_object@@PEAU1@@Z
+ ?TFE_TensorHandleToNumpy@tensorflow@@YAPEAU_object@@PEAUTFE_TensorHandle@@PEAUTSL_Status@@@Z
+ ?TF_CreatePlaceholders@tensorflow@@YA?AV?$vector@UTF_Output@@V?$allocator@UTF_Output@@@std@@@std@@PEAUTF_Graph@@PEAU_object@@PEBDPEAUTSL_Status@@@Z
+ ?TF_GraphGetTensorShapeHelper@tensorflow@@YA?AV?$InlinedVector@_J$05V?$allocator@_J@std@@@lts_20260526@absl@@PEAUTF_Graph@@UTF_Output@@PEAUTSL_Status@@PEA_N@Z
+ ?TF_GraphSetOutputHandleShapesAndTypes_wrapper@tensorflow@@YAXPEAUTF_Graph@@UTF_Output@@AEBV?$vector@V?$vector@_JV?$allocator@_J@std@@@std@@V?$allocator@V?$vector@_JV?$allocator@_J@std@@@std@@@2@@std@@AEBV?$vector@HV?$allocator@H@std@@@5@AEBV?$vector@W4TF_DataType@@V?$allocator@W4TF_DataType@@@std@@@5@PEAUTSL_Status@@@Z
+ ?TF_GraphSetTensorShape_wrapper@tensorflow@@YAXPEAUTF_Graph@@UTF_Output@@AEBV?$vector@_JV?$allocator@_J@std@@@std@@_NPEAUTSL_Status@@@Z
+ ?TF_GraphToFunction_wrapper@tensorflow@@YAPEAUTF_Function@@PEBUTF_Graph@@PEBD_NPEBV?$vector@PEAUTF_Operation@@V?$allocator@PEAUTF_Operation@@@std@@@std@@AEBV?$vector@UTF_Output@@V?$allocator@UTF_Output@@@std@@@5@4AEBV?$InlinedVector@PEBD$07V?$allocator@PEBD@std@@@lts_20260526@absl@@35PEBUTF_FunctionOptions@@1PEAUTSL_Status@@@Z
+ ?TF_ImportGraphDefResultsMissingUnusedInputMappings_wrapper@tensorflow@@YA?AV?$vector@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$allocator@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@@std@@PEAUTF_ImportGraphDefResults@@@Z
+ ?TF_NewSessionRef@tensorflow@@YAPEAUTF_Session@@PEAUTF_Graph@@PEBUTF_SessionOptions@@PEAUTSL_Status@@@Z
+ ?TF_OperationGetControlOutputs_wrapper@tensorflow@@YA?AV?$vector@PEAUTF_Operation@@V?$allocator@PEAUTF_Operation@@@std@@@std@@PEAUTF_Operation@@@Z
+ ?TF_Reset_wrapper@tensorflow@@YAXPEBUTF_SessionOptions@@AEBV?$InlinedVector@PEBD$07V?$allocator@PEBD@std@@@lts_20260526@absl@@PEAUTSL_Status@@@Z
+ ?TF_SessionMakeCallable@tensorflow@@YAXPEAUTF_Session@@PEBUTF_Buffer@@PEA_JPEAUTSL_Status@@@Z
+ ?TF_SessionPRunSetup_wrapper@tensorflow@@YAXPEAUTF_Session@@AEBV?$vector@UTF_Output@@V?$allocator@UTF_Output@@@std@@@std@@1AEBV?$vector@PEAUTF_Operation@@V?$allocator@PEAUTF_Operation@@@std@@@4@PEAPEBDPEAUTSL_Status@@@Z
+ ?TF_SessionPRun_wrapper@tensorflow@@YAXPEAUTF_Session@@PEBDAEBV?$vector@UTF_Output@@V?$allocator@UTF_Output@@@std@@@std@@AEBV?$vector@PEAU_object@@V?$allocator@PEAU_object@@@std@@@4@2PEAUTSL_Status@@PEAV54@@Z
+ ?TF_SessionReleaseCallable@tensorflow@@YAXPEAUTF_Session@@_JPEAUTSL_Status@@@Z
+ ?TF_SessionRunCallable@tensorflow@@YAXPEAUTF_Session@@_JPEAU_object@@PEAV?$InlinedVector@PEAU_object@@$07V?$allocator@PEAU_object@@@std@@@lts_20260526@absl@@PEAUTF_Buffer@@PEAUTSL_Status@@@Z
+ ?TF_SessionRun_wrapper@tensorflow@@YAXPEAUTF_Session@@PEBUTF_Buffer@@AEBV?$vector@UTF_Output@@V?$allocator@UTF_Output@@@std@@@std@@AEBV?$vector@PEAU_object@@V?$allocator@PEAU_object@@@std@@@5@2AEBV?$vector@PEAUTF_Operation@@V?$allocator@PEAUTF_Operation@@@std@@@5@PEAU3@PEAUTSL_Status@@PEAV65@@Z
+ ?TF_TryEvaluateConstant_wrapper@tensorflow@@YAPEAU_object@@PEAUTF_Graph@@UTF_Output@@PEAUTSL_Status@@@Z
+ ?TPUCoreIDsToLocations@dtensor@tensorflow@@YA?AV?$vector@V?$vector@HV?$allocator@H@std@@@std@@V?$allocator@V?$vector@HV?$allocator@H@std@@@std@@@2@@std@@PEAUTFE_Context@@AEBV?$vector@HV?$allocator@H@std@@@4@PEAX@Z
+ ?TPUCoreLocationsToIDs@dtensor@tensorflow@@YA?AV?$vector@HV?$allocator@H@std@@@std@@PEAUTFE_Context@@AEBV?$vector@V?$vector@HV?$allocator@H@std@@@std@@V?$allocator@V?$vector@HV?$allocator@H@std@@@std@@@2@@4@PEAX@Z
+ ?TensorToNdarray@tensorflow@@YA?AVStatus@lts_20260526@absl@@AEBVTensor@1@PEAPEAU_object@@@Z
+ ?TestCounter@metrics@tensorflow@@YAPEBVCounterCell@monitoring@tsl@@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@0@Z
+ ?ThrowBadStatusOrAccess@internal_statusor@lts_20260526@absl@@YAXVStatus@23@@Z
+ ?ThrowStdOutOfRange@lts_20260526@absl@@YAXPEBD@Z
+ ?ToDeviceType@Mesh@dtensor@tensorflow@@QEBA?AV?$StatusOr@VMesh@dtensor@tensorflow@@@lts_20260526@absl@@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@Z
+ ?ToFrames@FrozenStackTrace@tensorflow@@UEBA?AV?$Span@$$CBUStackFrame@tsl@@@lts_20260526@absl@@XZ
+ ?ToGraphDebugInfoStr@GraphDebugInfoBuilder@tensorflow@@QEBA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@XZ
+ ?ToGraphDef@Graph@tensorflow@@QEBAXPEAVGraphDef@2@_N1@Z
+ ?ToProto@CompileOptions@xla@@QEBA?AV?$StatusOr@VCompileOptionsProto@xla@@@lts_20260526@absl@@XZ
+ ?ToProto@Layout@dtensor@tensorflow@@QEBA?AV?$StatusOr@VLayoutProto@dtensor@tensorflow@@@lts_20260526@absl@@XZ
+ ?ToProto@Mesh@dtensor@tensorflow@@QEBA?AV?$StatusOr@VMeshProto@dtensor@tensorflow@@@lts_20260526@absl@@XZ
+ ?ToString@FrozenStackTrace@tensorflow@@UEBA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@AEBUTracePrintingOptions@AbstractStackTrace@2@@Z
+ ?ToString@HloInstruction@xla@@QEBA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@XZ
+ ?ToString@Layout@dtensor@tensorflow@@QEBA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@XZ
+ ?ToString@Mesh@dtensor@tensorflow@@QEBA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@XZ
+ ?ToStringSlow@Status@lts_20260526@absl@@CA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@_KW4StatusToStringMode@23@@Z
+ ?ToUncachedFrames@FrozenStackTrace@tensorflow@@UEBA?AV?$vector@UStackFrame@tsl@@V?$allocator@UStackFrame@tsl@@@std@@@std@@XZ
+ ?TocoConvert@toco@@YAPEAU_object@@PEAU2@00_N@Z
+ ?TopologicalSort@grappler@tensorflow@@YA?AVStatus@lts_20260526@absl@@PEAVGraphDef@2@@Z
+ ?TraceString@OpKernel@tensorflow@@UEBA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@AEBVOpKernelContext@2@_N@Z
+ ?TrainingTimeSaved@metrics@tensorflow@@YAAEAVCounterCell@monitoring@tsl@@V?$basic_string_view@DU?$char_traits@D@std@@@std@@@Z
+ ?TransformGraph@graph_transforms@tensorflow@@YA?AVStatus@lts_20260526@absl@@AEBV?$vector@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$allocator@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@@std@@0AEBV?$vector@U?$pair@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$map@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$vector@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$allocator@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@@2@U?$less@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@V?$allocator@U?$pair@$$CBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$vector@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$allocator@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@@2@@std@@@2@@2@@std@@V?$allocator@U?$pair@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$map@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$vector@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$allocator@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@@2@U?$less@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@V?$allocator@U?$pair@$$CBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$vector@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$allocator@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@@2@@std@@@2@@2@@std@@@2@@7@PEAVGraphDef@2@@Z
+ ?TreatAsSet@MessageDifferencer@util@protobuf@google@@QEAAXPEBVFieldDescriptor@34@@Z
+ ?TryFindKernelClass@swig@tensorflow@@YA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@AEBV34@@Z
+ ?TypeConstraint@KernelDefBuilder@tensorflow@@QEAAAEAV12@PEBDW4DataType@2@@Z
+ ?UnchangedShape@shape_inference@tensorflow@@YA?AVStatus@lts_20260526@absl@@PEAVInferenceContext@12@@Z
+ ?UnknownShape@InferenceContext@shape_inference@tensorflow@@QEAA?AVShapeHandle@23@XZ
+ ?UnknownShape@shape_inference@tensorflow@@YA?AVStatus@lts_20260526@absl@@PEAVInferenceContext@12@@Z
+ ?UnknownShapeOfRank@InferenceContext@shape_inference@tensorflow@@QEAA?AVShapeHandle@23@_J@Z
+ ?Unpack@dtensor@tensorflow@@YA?AV?$vector@PEAUTFE_TensorHandle@@V?$allocator@PEAUTFE_TensorHandle@@@std@@@std@@PEAUTFE_Context@@PEAUTFE_TensorHandle@@PEAXPEAUTSL_Status@@@Z
+ ?Unref@StatusRep@status_internal@lts_20260526@absl@@QEBAXXZ
+ ?Unregister@CollectionRegistry@monitoring@tsl@@AEAAXPEBVAbstractMetricDef@23@@Z
+ ?Unregister@PythonAPIDispatcher@py_dispatch@tensorflow@@QEAAXPEAU_object@@@Z
+ ?UpdateEdge@tensorflow@@YAXPEAUTF_Graph@@UTF_Output@@UTF_Input@@PEAUTSL_Status@@@Z
+ ?ValidateType@ResourceHandle@tensorflow@@QEBA?AVStatus@lts_20260526@absl@@AEBVTypeIndex@2@@Z
+ ?Vector@InferenceContext@shape_inference@tensorflow@@QEAA?AVShapeHandle@23@UDimensionOrConstant@23@@Z
+ ?WaitCommon@CondVar@lts_20260526@absl@@AEAA_NPEAVMutex@23@VKernelTimeout@synchronization_internal@23@@Z
+ ?Watch@Tape@gradients@tensorflow@@QEAAXPEBVAbstractTensorHandle@3@@Z
+ ?WithRank@InferenceContext@shape_inference@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@VShapeHandle@23@_JPEAV723@@Z
+ ?WithVerbosity@LogMessage@log_internal@lts_20260526@absl@@QEAAAEAV1234@H@Z
+ ?WriteProfile@tfprof@tensorflow@@YAXPEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@Z
+ ?WriteRecord@RecordWriter@io@tsl@@QEAA?AVStatus@lts_20260526@absl@@V?$basic_string_view@DU?$char_traits@D@std@@@std@@@Z
+ ?WriteSerializedEvent@EventsWriter@tensorflow@@QEAAXV?$basic_string_view@DU?$char_traits@D@std@@@std@@@Z
+ ?WriteSerializedExecutionDebugEvent@DebugEventsWriter@tfdbg@tensorflow@@QEAAXAEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@W4DebugEventFileType@23@@Z
+ ?WriteSerializedNonExecutionDebugEvent@DebugEventsWriter@tfdbg@tensorflow@@QEAAXAEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@W4DebugEventFileType@23@@Z
+ ?WriteStringToFile@tsl@@YA?AVStatus@lts_20260526@absl@@PEAVEnv@1@AEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$basic_string_view@DU?$char_traits@D@std@@@7@@Z
+ ?_AttrValue_ListValue_default_instance_@tensorflow@@3UAttrValue_ListValueDefaultTypeInternal@1@A
+ ?_FunctionDefLibrary_default_instance_@tensorflow@@3UFunctionDefLibraryDefaultTypeInternal@1@A
+ ?_FunctionDef_default_instance_@tensorflow@@3UFunctionDefDefaultTypeInternal@1@A
+ ?_GraphDef_default_instance_@tensorflow@@3UGraphDefDefaultTypeInternal@1@A
+ ?_LayoutProto_default_instance_@dtensor@tensorflow@@3ULayoutProtoDefaultTypeInternal@12@A
+ ?_MeshProto_default_instance_@dtensor@tensorflow@@3UMeshProtoDefaultTypeInternal@12@A
+ ?_MetaGraphDef_default_instance_@tensorflow@@3UMetaGraphDefDefaultTypeInternal@1@A
+ ?add_ret_check_failure@MakeErrorStream@status_macros@xla@@QEAAAEAVMakeErrorStreamWithOutput@123@PEBD@Z
+ ?checkInput@FileCheck@llvm@@QEAA_NAEAVSourceMgr@2@VStringRef@2@PEAVFileCheckDiagList@2@@Z
+ ?clear_value@AttrValue@tensorflow@@QEAAXXZ
+ ?cost@PyListChecker@py_dispatch@tensorflow@@UEBAHXZ
+ ?cost@PyUnionChecker@py_dispatch@tensorflow@@UEBAHXZ
+ ?create@ModuleOp@mlir@@SA?AV12@AEAVOpBuilder@2@VLocation@2@V?$optional@VStringRef@llvm@@@std@@@Z
+ ?deallocate_buffer@llvm@@YAXPEAX_K1@Z
+ ?def@Node@tensorflow@@QEBAAEBVNodeDef@2@XZ
+ ?device_location@Mesh@dtensor@tensorflow@@QEBA?AV?$StatusOr@$$CBV?$InlinedVector@_J$03V?$allocator@_J@std@@@lts_20260526@absl@@@lts_20260526@absl@@H@Z
+ ?device_type@Mesh@dtensor@tensorflow@@QEBA?AV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@XZ
+ ?dim_size@?$TensorShapeBase@VPartialTensorShape@tensorflow@@@tensorflow@@QEBA_JH@Z
+ ?dim_size@?$TensorShapeBase@VTensorShape@tensorflow@@@tensorflow@@QEBA_JH@Z
+ ?dim_size@Mesh@dtensor@tensorflow@@QEBA?AV?$StatusOr@_J@lts_20260526@absl@@V?$basic_string_view@DU?$char_traits@D@std@@@std@@@Z
+ ?dim_sizes@?$TensorShapeBase@VPartialTensorShape@tensorflow@@@tensorflow@@QEBA?AV?$InlinedVector@_J$03V?$allocator@_J@std@@@lts_20260526@absl@@XZ
+ ?dim_sizes@Mesh@dtensor@tensorflow@@QEBA?AV?$vector@_JV?$allocator@_J@std@@@std@@XZ
+ ?emitError@OpState@mlir@@QEAA?AVInFlightDiagnostic@2@AEBVTwine@llvm@@@Z
+ ?enable_tensor_float_32_execution@tsl@@YAX_N@Z
+ ?erase@Operation@mlir@@QEAAXXZ
+ ?fixed_address_empty_string@internal@protobuf@google@@3VGlobalEmptyStringDynamicInit@123@A
+ ?getMemBuffer@MemoryBuffer@llvm@@SA?AV?$unique_ptr@VMemoryBuffer@llvm@@U?$default_delete@VMemoryBuffer@llvm@@@std@@@std@@VStringRef@2@0_N@Z
+ ?getOrLoadDialect@MLIRContext@mlir@@QEAAPEAVDialect@2@VStringRef@llvm@@VTypeID@2@V?$function_ref@$$A6A?AV?$unique_ptr@VDialect@mlir@@U?$default_delete@VDialect@mlir@@@std@@@std@@XZ@5@@Z
+ ?getValuePtr@?$SpecificNodeAccess@U?$node_options@VOperation@mlir@@$0A@$0A@X$0A@X@ilist_detail@llvm@@@ilist_detail@llvm@@KAPEAVOperation@mlir@@PEAV?$ilist_node_impl@U?$node_options@VOperation@mlir@@$0A@$0A@X$0A@X@ilist_detail@llvm@@@3@@Z
+ ?id@?$TypeIDResolver@VArithDialect@arith@mlir@@X@detail@mlir@@0VSelfOwningTypeID@3@A
+ ?id@?$TypeIDResolver@VFuncDialect@func@mlir@@X@detail@mlir@@0VSelfOwningTypeID@3@A
+ ?id@?$TypeIDResolver@VModuleOp@mlir@@X@detail@mlir@@0VSelfOwningTypeID@3@A
+ ?id@?$TypeIDResolver@VSCFDialect@scf@mlir@@X@detail@mlir@@0VSelfOwningTypeID@3@A
+ ?id@?$TypeIDResolver@VShapeDialect@shape@mlir@@X@detail@mlir@@0VSelfOwningTypeID@3@A
+ ?input@InferenceContext@shape_inference@tensorflow@@QEBA?AVStatus@lts_20260526@absl@@V?$basic_string_view@DU?$char_traits@D@std@@@std@@PEAV?$vector@VShapeHandle@shape_inference@tensorflow@@V?$allocator@VShapeHandle@shape_inference@tensorflow@@@std@@@8@@Z
+ ?input@OpKernelContext@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@V?$basic_string_view@DU?$char_traits@D@std@@@std@@PEAPEBVTensor@2@@Z
+ ?input@OpKernelContext@tensorflow@@QEBAAEBVTensor@2@H@Z
+ ?insert@DialectRegistry@mlir@@QEAAXVTypeID@2@VStringRef@llvm@@AEBV?$function@$$A6APEAVDialect@mlir@@PEAVMLIRContext@2@@Z@std@@@Z
+ ?internal_log_function@raw_log_internal@lts_20260526@absl@@3V?$AtomicHook@P6AXW4LogSeverity@lts_20260526@absl@@PEBDHAEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@Z@base_internal@23@A
+ ?loadAllAvailableDialects@MLIRContext@mlir@@QEAAXXZ
+ ?lock@Mutex@lts_20260526@absl@@QEAAXXZ
+ ?make_safe@tensorflow@@YA?AV?$unique_ptr@UTF_Buffer@@UTFBufferDeleter@detail@tensorflow@@@std@@PEAUTF_Buffer@@@Z
+ ?make_safe@tensorflow@@YA?AV?$unique_ptr@UTSL_Status@@UTFStatusDeleter@detail@tensorflow@@@std@@PEAUTSL_Status@@@Z
+ ?make_safe@tensorflow@@YA?AV?$unique_ptr@U_object@@UPyDecrefDeleter@detail@tensorflow@@@std@@PEAU_object@@@Z
+ ?name@Node@tensorflow@@QEBAAEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@XZ
+ ?num_devices@Mesh@dtensor@tensorflow@@QEBA_JXZ
+ ?num_shards_for_dim@Layout@dtensor@tensorflow@@QEBA_KH@Z
+ ?op_def@Node@tensorflow@@QEBAAEBVOpDef@2@XZ
+ ?parseSourceFile@mlir@@YA?AULogicalResult@llvm@@AEBVSourceMgr@3@PEAVBlock@1@AEBVParserConfig@1@PEAVLocationAttr@1@@Z
+ ?pywrap_library_dependency_symbol@python@tensorflow@@YAXXZ
+ ?readCheckFile@FileCheck@llvm@@QEAA_NAEAVSourceMgr@2@VStringRef@2@PEAU?$pair@II@std@@@Z
+ ?record_persistent_memory_allocation@OpKernelContext@tensorflow@@QEAAX_J0@Z
+ ?remove@Operation@mlir@@QEAAXXZ
+ ?report@InFlightDiagnostic@mlir@@QEAAXXZ
+ ?requested_device@Node@tensorflow@@QEBAAEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@XZ
+ ?set_field_comparator@MessageDifferencer@util@protobuf@google@@QEAAXPEAVDefaultFieldComparator@234@@Z
+ ?set_output@InferenceContext@shape_inference@tensorflow@@QEAA?AVStatus@lts_20260526@absl@@V?$basic_string_view@DU?$char_traits@D@std@@@std@@AEBV?$vector@VShapeHandle@shape_inference@tensorflow@@V?$allocator@VShapeHandle@shape_inference@tensorflow@@@std@@@8@@Z
+ ?set_output@OpKernelContext@tensorflow@@QEAAXHAEBVTensor@2@@Z
+ ?set_requested_device@Node@tensorflow@@QEAAXAEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@Z
+ ?set_tf2_execution@tensorflow@@YAX_N@Z
+ ?sharding_spec_strs@Layout@dtensor@tensorflow@@QEBA?AV?$vector@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@V?$allocator@V?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@@2@@std@@XZ
+ ?stream@OstreamView@LogMessage@log_internal@lts_20260526@absl@@QEAAAEAV?$basic_ostream@DU?$char_traits@D@std@@@std@@XZ
+ ?tensor_data@Tensor@tensorflow@@QEBA?AV?$basic_string_view@DU?$char_traits@D@std@@@std@@XZ
+ ?tensor_float_32_execution_enabled@tsl@@YA_NXZ
+ ?tf2_execution_enabled@tensorflow@@YA_NXZ
+ ?transferNodesFromList@?$ilist_traits@VOperation@mlir@@@llvm@@QEAAXAEAU12@V?$ilist_iterator@U?$node_options@VOperation@mlir@@$0A@$0A@X$0A@X@ilist_detail@llvm@@$0A@$0A@@2@1@Z
+ ?type_string@Node@tensorflow@@QEBAAEBV?$basic_string@DU?$char_traits@D@std@@V?$allocator@D@2@@std@@XZ
+ ?unlock@Mutex@lts_20260526@absl@@QEAAXXZ
+ ?verify@mlir@@YA?AULogicalResult@llvm@@PEAVOperation@1@_N@Z
+ ?verifyInvariants@ModuleOp@mlir@@QEAA?AULogicalResult@llvm@@XZ
+ ?versions@Graph@tensorflow@@QEBAAEBVVersionDef@2@XZ
+ TF_CreateFunction
+ TF_SetTracingImplementation
+; go/keep-sorted end

tensorflow/python/_pywrap_tfcompile.pyi

@@ -0,0 +1,16 @@
+# 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.
+# ==============================================================================
+
+def Compile(graph: str = ..., config: str = ..., target_triple: str = ..., target_cpu: str = ..., target_features: str = ..., entry_point: str = ..., cpp_class: str = ..., out_function_object: str = ..., out_metadata_object: str = ..., out_header: str = ..., out_session_module: str = ..., out_constant_buffers_object: str = ..., mlir_components: str = ..., gen_name_to_index: bool = ..., gen_program_shape: bool = ...) -> None: ...

tensorflow/python/_pywrap_tfe.pyi

@@ -0,0 +1,371 @@
+# 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 typing import ClassVar
+
+TFE_DEVICE_PLACEMENT_EXPLICIT: TFE_ContextDevicePlacementPolicy
+TFE_DEVICE_PLACEMENT_SILENT: TFE_ContextDevicePlacementPolicy
+TFE_DEVICE_PLACEMENT_SILENT_FOR_INT32: TFE_ContextDevicePlacementPolicy
+TFE_DEVICE_PLACEMENT_WARN: TFE_ContextDevicePlacementPolicy
+TF_ATTR_BOOL: TF_AttrType
+TF_ATTR_FLOAT: TF_AttrType
+TF_ATTR_FUNC: TF_AttrType
+TF_ATTR_INT: TF_AttrType
+TF_ATTR_PLACEHOLDER: TF_AttrType
+TF_ATTR_SHAPE: TF_AttrType
+TF_ATTR_STRING: TF_AttrType
+TF_ATTR_TENSOR: TF_AttrType
+TF_ATTR_TYPE: TF_AttrType
+
+class EagerContextThreadLocalData:
+    device_name: object
+    device_spec: object
+    executor: object
+    function_call_options: object
+    invoking_op_callbacks: bool
+    is_eager: bool
+    op_callbacks: object
+    scope_name: object
+    def __init__(self, py_eager_context: object, is_eager: object, device_spec: object) -> None: ...
+
+class TFE_CancellationManager:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class TFE_ContextDevicePlacementPolicy:
+    __members__: ClassVar[dict] = ...  # read-only
+    TFE_DEVICE_PLACEMENT_EXPLICIT: ClassVar[TFE_ContextDevicePlacementPolicy] = ...
+    TFE_DEVICE_PLACEMENT_SILENT: ClassVar[TFE_ContextDevicePlacementPolicy] = ...
+    TFE_DEVICE_PLACEMENT_SILENT_FOR_INT32: ClassVar[TFE_ContextDevicePlacementPolicy] = ...
+    TFE_DEVICE_PLACEMENT_WARN: ClassVar[TFE_ContextDevicePlacementPolicy] = ...
+    __entries: ClassVar[dict] = ...
+    def __init__(self, value: int) -> None: ...
+    def __eq__(self, other: object) -> bool: ...
+    def __hash__(self) -> int: ...
+    def __index__(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __ne__(self, other: object) -> bool: ...
+    @property
+    def name(self) -> str: ...
+    @property
+    def value(self) -> int: ...
+
+class TFE_ContextOptions:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class TFE_Executor:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class TFE_MonitoringBoolGauge0:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class TFE_MonitoringBoolGauge1:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class TFE_MonitoringBoolGauge2:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class TFE_MonitoringBoolGaugeCell:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class TFE_MonitoringBuckets:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class TFE_MonitoringCounter0:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class TFE_MonitoringCounter1:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class TFE_MonitoringCounter2:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class TFE_MonitoringCounterCell:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class TFE_MonitoringIntGauge0:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class TFE_MonitoringIntGauge1:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class TFE_MonitoringIntGauge2:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class TFE_MonitoringIntGaugeCell:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class TFE_MonitoringSampler0:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class TFE_MonitoringSampler1:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class TFE_MonitoringSampler2:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class TFE_MonitoringSamplerCell:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class TFE_MonitoringStringGauge0:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class TFE_MonitoringStringGauge1:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class TFE_MonitoringStringGauge2:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class TFE_MonitoringStringGauge3:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class TFE_MonitoringStringGauge4:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class TFE_MonitoringStringGaugeCell:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class TF_AttrType:
+    __members__: ClassVar[dict] = ...  # read-only
+    TF_ATTR_BOOL: ClassVar[TF_AttrType] = ...
+    TF_ATTR_FLOAT: ClassVar[TF_AttrType] = ...
+    TF_ATTR_FUNC: ClassVar[TF_AttrType] = ...
+    TF_ATTR_INT: ClassVar[TF_AttrType] = ...
+    TF_ATTR_PLACEHOLDER: ClassVar[TF_AttrType] = ...
+    TF_ATTR_SHAPE: ClassVar[TF_AttrType] = ...
+    TF_ATTR_STRING: ClassVar[TF_AttrType] = ...
+    TF_ATTR_TENSOR: ClassVar[TF_AttrType] = ...
+    TF_ATTR_TYPE: ClassVar[TF_AttrType] = ...
+    __entries: ClassVar[dict] = ...
+    def __init__(self, value: int) -> None: ...
+    def __eq__(self, other: object) -> bool: ...
+    def __hash__(self) -> int: ...
+    def __index__(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __ne__(self, other: object) -> bool: ...
+    @property
+    def name(self) -> str: ...
+    @property
+    def value(self) -> int: ...
+
+class TF_Buffer:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class TF_DeviceList:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+class TF_Function:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+def TFE_AbortCollectiveOps(arg0: object, arg1: int, arg2: str) -> None: ...
+def TFE_CancellationManagerIsCancelled(arg0: TFE_CancellationManager) -> bool: ...
+def TFE_CancellationManagerStartCancel(arg0: TFE_CancellationManager) -> None: ...
+def TFE_ClearScalarCache() -> object: ...
+def TFE_CollectiveOpsCheckPeerHealth(arg0: object, arg1: str, arg2: int) -> None: ...
+def TFE_ContextAddFunction(arg0: object, arg1: TF_Function) -> None: ...
+def TFE_ContextAddFunctionDef(arg0: object, arg1: str, arg2: int) -> None: ...
+def TFE_ContextAddFunctionDefNoSerialization(ctx: object, function_def) -> None: ...
+def TFE_ContextCheckAlive(arg0: object, arg1: str) -> bool: ...
+def TFE_ContextClearCaches(arg0: object) -> None: ...
+def TFE_ContextClearExecutors(arg0: object) -> None: ...
+def TFE_ContextDisableGraphCollection(arg0: object) -> None: ...
+def TFE_ContextDisableRunMetadata(arg0: object) -> None: ...
+def TFE_ContextEnableGraphCollection(arg0: object) -> None: ...
+def TFE_ContextEnableRunMetadata(arg0: object) -> None: ...
+def TFE_ContextExportRunMetadata(arg0: object, arg1: TF_Buffer) -> None: ...
+def TFE_ContextGetDevicePlacementPolicy(arg0: object) -> TFE_ContextDevicePlacementPolicy: ...
+def TFE_ContextGetExecutorForThread(arg0: object) -> TFE_Executor: ...
+def TFE_ContextGetFunction(arg0: object, arg1: str) -> TF_Function: ...
+def TFE_ContextGetFunctionDef(arg0: object, arg1: str, arg2: TF_Buffer) -> None: ...
+def TFE_ContextGetFunctionDefNoSerialization(*args, **kwargs): ...
+def TFE_ContextGetGraphDebugInfo(arg0: object, arg1: str, arg2: TF_Buffer) -> None: ...
+def TFE_ContextHasFunction(arg0: object, arg1: str) -> int: ...
+def TFE_ContextListDevices(arg0: object) -> TF_DeviceList: ...
+def TFE_ContextListFunctionNames(arg0: object) -> list[str]: ...
+def TFE_ContextOptionsSetAsync(arg0: TFE_ContextOptions, arg1: int) -> None: ...
+def TFE_ContextOptionsSetConfig(arg0: TFE_ContextOptions, arg1: bytes) -> None: ...
+def TFE_ContextOptionsSetDevicePlacementPolicy(arg0: TFE_ContextOptions, arg1: TFE_ContextDevicePlacementPolicy) -> None: ...
+def TFE_ContextOptionsSetJitCompileRewrite(arg0: TFE_ContextOptions, arg1: bool) -> None: ...
+def TFE_ContextOptionsSetRunEagerOpAsFunction(arg0: TFE_ContextOptions, arg1: bool) -> None: ...
+def TFE_ContextOptionsSetTfrt(arg0: TFE_ContextOptions, arg1: bool) -> None: ...
+def TFE_ContextRemoveFunction(arg0: object, arg1: str) -> None: ...
+def TFE_ContextSetExecutorForThread(arg0: object, arg1: TFE_Executor) -> None: ...
+def TFE_ContextSetJitCompileRewrite(arg0: object, arg1: bool) -> None: ...
+def TFE_ContextSetLogDevicePlacement(arg0: object, arg1: bool) -> None: ...
+def TFE_ContextSetRunEagerOpAsFunction(arg0: object, arg1: bool) -> None: ...
+def TFE_ContextSetServerDef(arg0: object, arg1: int, arg2: bytes) -> None: ...
+def TFE_ContextSetServerDefWithTimeoutAndRetries(arg0: object, arg1: int, arg2: bytes, arg3: int, arg4: int) -> None: ...
+def TFE_ContextSetSoftDevicePlacement(arg0: object, arg1: bool) -> None: ...
+def TFE_ContextSetThreadLocalDevicePlacementPolicy(arg0: object, arg1: TFE_ContextDevicePlacementPolicy) -> None: ...
+def TFE_ContextSyncExecutors(arg0: object) -> None: ...
+def TFE_ContextUpdateServerDef(arg0: object, arg1: int, arg2: bytes) -> None: ...
+def TFE_DeleteConfigKeyValue(arg0: object, arg1: str) -> None: ...
+def TFE_DeleteContext(arg0: object) -> None: ...
+def TFE_DeleteContextOptions(arg0: TFE_ContextOptions) -> None: ...
+def TFE_DeleteExecutor(arg0: TFE_Executor) -> None: ...
+def TFE_DlpackDevice(arg0: object) -> tuple: ...
+def TFE_EnableCollectiveOps(arg0: object, arg1: bytes) -> None: ...
+def TFE_ExecutorClearError(arg0: TFE_Executor) -> None: ...
+def TFE_ExecutorIsAsync(arg0: TFE_Executor) -> bool: ...
+def TFE_ExecutorWaitForAllPendingNodes(arg0: TFE_Executor) -> None: ...
+def TFE_FromDlpackCapsule(arg0, arg1: object) -> object: ...
+def TFE_GetConfigKeyValue(arg0: object, arg1: str, arg2: int, arg3: TF_Buffer) -> None: ...
+def TFE_GetContextId(arg0: object) -> int: ...
+def TFE_GetMemoryInfo(arg0: object, arg1: str) -> dict[str, int]: ...
+def TFE_GetTaskStates(arg0: object, arg1: list[str], arg2: list[int]) -> object: ...
+def TFE_HostAddressSpace(arg0: object, arg1: TF_Buffer) -> None: ...
+def TFE_InsertConfigKeyValue(arg0: object, arg1: str, arg2: str) -> None: ...
+def TFE_MonitoringBoolGaugeCellSet(arg0: TFE_MonitoringBoolGaugeCell, arg1: bool) -> None: ...
+def TFE_MonitoringBoolGaugeCellValue(arg0: TFE_MonitoringBoolGaugeCell) -> bool: ...
+def TFE_MonitoringCounterCellIncrementBy(arg0: TFE_MonitoringCounterCell, arg1: int) -> None: ...
+def TFE_MonitoringCounterCellValue(arg0: TFE_MonitoringCounterCell) -> int: ...
+def TFE_MonitoringDeleteBoolGauge0(arg0: TFE_MonitoringBoolGauge0) -> None: ...
+def TFE_MonitoringDeleteBoolGauge1(arg0: TFE_MonitoringBoolGauge1) -> None: ...
+def TFE_MonitoringDeleteBoolGauge2(arg0: TFE_MonitoringBoolGauge2) -> None: ...
+def TFE_MonitoringDeleteBuckets(arg0: TFE_MonitoringBuckets) -> None: ...
+def TFE_MonitoringDeleteCounter0(arg0: TFE_MonitoringCounter0) -> None: ...
+def TFE_MonitoringDeleteCounter1(arg0: TFE_MonitoringCounter1) -> None: ...
+def TFE_MonitoringDeleteCounter2(arg0: TFE_MonitoringCounter2) -> None: ...
+def TFE_MonitoringDeleteIntGauge0(arg0: TFE_MonitoringIntGauge0) -> None: ...
+def TFE_MonitoringDeleteIntGauge1(arg0: TFE_MonitoringIntGauge1) -> None: ...
+def TFE_MonitoringDeleteIntGauge2(arg0: TFE_MonitoringIntGauge2) -> None: ...
+def TFE_MonitoringDeleteSampler0(arg0: TFE_MonitoringSampler0) -> None: ...
+def TFE_MonitoringDeleteSampler1(arg0: TFE_MonitoringSampler1) -> None: ...
+def TFE_MonitoringDeleteSampler2(arg0: TFE_MonitoringSampler2) -> None: ...
+def TFE_MonitoringDeleteStringGauge0(arg0: TFE_MonitoringStringGauge0) -> None: ...
+def TFE_MonitoringDeleteStringGauge1(arg0: TFE_MonitoringStringGauge1) -> None: ...
+def TFE_MonitoringDeleteStringGauge2(arg0: TFE_MonitoringStringGauge2) -> None: ...
+def TFE_MonitoringDeleteStringGauge3(arg0: TFE_MonitoringStringGauge3) -> None: ...
+def TFE_MonitoringDeleteStringGauge4(arg0: TFE_MonitoringStringGauge4) -> None: ...
+def TFE_MonitoringGetCellBoolGauge0(arg0: TFE_MonitoringBoolGauge0) -> TFE_MonitoringBoolGaugeCell: ...
+def TFE_MonitoringGetCellBoolGauge1(arg0: TFE_MonitoringBoolGauge1, arg1: str) -> TFE_MonitoringBoolGaugeCell: ...
+def TFE_MonitoringGetCellBoolGauge2(arg0: TFE_MonitoringBoolGauge2, arg1: str, arg2: str) -> TFE_MonitoringBoolGaugeCell: ...
+def TFE_MonitoringGetCellCounter0(arg0: TFE_MonitoringCounter0) -> TFE_MonitoringCounterCell: ...
+def TFE_MonitoringGetCellCounter1(arg0: TFE_MonitoringCounter1, arg1: str) -> TFE_MonitoringCounterCell: ...
+def TFE_MonitoringGetCellCounter2(arg0: TFE_MonitoringCounter2, arg1: str, arg2: str) -> TFE_MonitoringCounterCell: ...
+def TFE_MonitoringGetCellIntGauge0(arg0: TFE_MonitoringIntGauge0) -> TFE_MonitoringIntGaugeCell: ...
+def TFE_MonitoringGetCellIntGauge1(arg0: TFE_MonitoringIntGauge1, arg1: str) -> TFE_MonitoringIntGaugeCell: ...
+def TFE_MonitoringGetCellIntGauge2(arg0: TFE_MonitoringIntGauge2, arg1: str, arg2: str) -> TFE_MonitoringIntGaugeCell: ...
+def TFE_MonitoringGetCellSampler0(arg0: TFE_MonitoringSampler0) -> TFE_MonitoringSamplerCell: ...
+def TFE_MonitoringGetCellSampler1(arg0: TFE_MonitoringSampler1, arg1: str) -> TFE_MonitoringSamplerCell: ...
+def TFE_MonitoringGetCellSampler2(arg0: TFE_MonitoringSampler2, arg1: str, arg2: str) -> TFE_MonitoringSamplerCell: ...
+def TFE_MonitoringGetCellStringGauge0(arg0: TFE_MonitoringStringGauge0) -> TFE_MonitoringStringGaugeCell: ...
+def TFE_MonitoringGetCellStringGauge1(arg0: TFE_MonitoringStringGauge1, arg1: str) -> TFE_MonitoringStringGaugeCell: ...
+def TFE_MonitoringGetCellStringGauge2(arg0: TFE_MonitoringStringGauge2, arg1: str, arg2: str) -> TFE_MonitoringStringGaugeCell: ...
+def TFE_MonitoringGetCellStringGauge3(arg0: TFE_MonitoringStringGauge3, arg1: str, arg2: str, arg3: str) -> TFE_MonitoringStringGaugeCell: ...
+def TFE_MonitoringGetCellStringGauge4(arg0: TFE_MonitoringStringGauge4, arg1: str, arg2: str, arg3: str, arg4: str) -> TFE_MonitoringStringGaugeCell: ...
+def TFE_MonitoringIntGaugeCellSet(arg0: TFE_MonitoringIntGaugeCell, arg1: int) -> None: ...
+def TFE_MonitoringIntGaugeCellValue(arg0: TFE_MonitoringIntGaugeCell) -> int: ...
+def TFE_MonitoringNewBoolGauge0(arg0: str, arg1: str) -> TFE_MonitoringBoolGauge0: ...
+def TFE_MonitoringNewBoolGauge1(arg0: str, arg1: str, arg2: str) -> TFE_MonitoringBoolGauge1: ...
+def TFE_MonitoringNewBoolGauge2(arg0: str, arg1: str, arg2: str, arg3: str) -> TFE_MonitoringBoolGauge2: ...
+def TFE_MonitoringNewCounter0(arg0: str, arg1: str) -> TFE_MonitoringCounter0: ...
+def TFE_MonitoringNewCounter1(arg0: str, arg1: str, arg2: str) -> TFE_MonitoringCounter1: ...
+def TFE_MonitoringNewCounter2(arg0: str, arg1: str, arg2: str, arg3: str) -> TFE_MonitoringCounter2: ...
+def TFE_MonitoringNewExponentialBuckets(arg0: float, arg1: float, arg2: int) -> TFE_MonitoringBuckets: ...
+def TFE_MonitoringNewIntGauge0(arg0: str, arg1: str) -> TFE_MonitoringIntGauge0: ...
+def TFE_MonitoringNewIntGauge1(arg0: str, arg1: str, arg2: str) -> TFE_MonitoringIntGauge1: ...
+def TFE_MonitoringNewIntGauge2(arg0: str, arg1: str, arg2: str, arg3: str) -> TFE_MonitoringIntGauge2: ...
+def TFE_MonitoringNewSampler0(arg0: str, arg1: TFE_MonitoringBuckets, arg2: str) -> TFE_MonitoringSampler0: ...
+def TFE_MonitoringNewSampler1(arg0: str, arg1: TFE_MonitoringBuckets, arg2: str, arg3: str) -> TFE_MonitoringSampler1: ...
+def TFE_MonitoringNewSampler2(arg0: str, arg1: TFE_MonitoringBuckets, arg2: str, arg3: str, arg4: str) -> TFE_MonitoringSampler2: ...
+def TFE_MonitoringNewStringGauge0(arg0: str, arg1: str) -> TFE_MonitoringStringGauge0: ...
+def TFE_MonitoringNewStringGauge1(arg0: str, arg1: str, arg2: str) -> TFE_MonitoringStringGauge1: ...
+def TFE_MonitoringNewStringGauge2(arg0: str, arg1: str, arg2: str, arg3: str) -> TFE_MonitoringStringGauge2: ...
+def TFE_MonitoringNewStringGauge3(arg0: str, arg1: str, arg2: str, arg3: str, arg4: str) -> TFE_MonitoringStringGauge3: ...
+def TFE_MonitoringNewStringGauge4(arg0: str, arg1: str, arg2: str, arg3: str, arg4: str, arg5: str) -> TFE_MonitoringStringGauge4: ...
+def TFE_MonitoringSamplerCellAdd(arg0: TFE_MonitoringSamplerCell, arg1: float) -> None: ...
+def TFE_MonitoringSamplerCellValue(arg0: TFE_MonitoringSamplerCell, arg1: TF_Buffer) -> None: ...
+def TFE_MonitoringStringGaugeCellSet(arg0: TFE_MonitoringStringGaugeCell, arg1: str) -> None: ...
+def TFE_MonitoringStringGaugeCellValue(arg0: TFE_MonitoringStringGaugeCell, arg1: TF_Buffer) -> None: ...
+def TFE_NewCancellationManager() -> TFE_CancellationManager: ...
+def TFE_NewContext(arg0: TFE_ContextOptions) -> object: ...
+def TFE_NewContextOptions() -> TFE_ContextOptions: ...
+def TFE_NewExecutor(arg0: bool, arg1: bool, arg2: int) -> TFE_Executor: ...
+def TFE_OpNameGetAttrType(arg0: object, arg1: str, arg2: str) -> object: ...
+def TFE_Py_EnableInteractivePythonLogging() -> None: ...
+def TFE_Py_Execute(arg0: object, arg1: str, arg2: str, arg3: object, arg4: object, arg5: object) -> object: ...
+def TFE_Py_ExecuteCancelable(arg0: object, arg1: str, arg2: str, arg3: object, arg4: object, arg5: TFE_CancellationManager, arg6: object) -> object: ...
+def TFE_Py_FastPathExecute(*args) -> object: ...
+def TFE_Py_ForwardAccumulatorJVP(arg0: object, arg1: object) -> object: ...
+def TFE_Py_ForwardAccumulatorNew(arg0: bool) -> object: ...
+def TFE_Py_ForwardAccumulatorPopState() -> object: ...
+def TFE_Py_ForwardAccumulatorPushState() -> object: ...
+def TFE_Py_ForwardAccumulatorSetAdd(arg0: object) -> object: ...
+def TFE_Py_ForwardAccumulatorSetRemove(arg0: object) -> None: ...
+def TFE_Py_ForwardAccumulatorWatch(arg0: object, arg1: object, arg2: object) -> None: ...
+def TFE_Py_InitEagerTensor(arg0: object) -> object: ...
+def TFE_Py_IsCustomDevice(arg0: object, arg1: str) -> bool: ...
+def TFE_Py_PackEagerTensors(arg0: object, arg1: object) -> object: ...
+def TFE_Py_PackJVPs(arg0: object) -> object: ...
+def TFE_Py_RecordGradient(arg0: object, arg1: object, arg2: object, arg3: object, arg4: object) -> object: ...
+def TFE_Py_RegisterCustomDevice(arg0: object, arg1, arg2: str, arg3) -> None: ...
+def TFE_Py_RegisterExceptionClass(arg0: object) -> object: ...
+def TFE_Py_RegisterFallbackExceptionClass(arg0: object) -> object: ...
+def TFE_Py_RegisterGradientFunction(arg0: object) -> object: ...
+def TFE_Py_RegisterJVPFunction(arg0: object) -> object: ...
+def TFE_Py_RegisterVSpace(arg0: object) -> object: ...
+def TFE_Py_SetCEagerContext(arg0: object) -> None: ...
+def TFE_Py_SetEagerContext(arg0: object) -> object: ...
+def TFE_Py_SetEagerTensorProfiler(arg0: object) -> object: ...
+def TFE_Py_TapeGradient(arg0: object, arg1: object, arg2: object, arg3: object, arg4: object, arg5: object) -> object: ...
+def TFE_Py_TapeSetAdd(arg0: object) -> None: ...
+def TFE_Py_TapeSetDeleteTrace(arg0: int) -> None: ...
+def TFE_Py_TapeSetIsEmpty() -> object: ...
+def TFE_Py_TapeSetIsStopped() -> object: ...
+def TFE_Py_TapeSetNew(arg0: object, arg1: object) -> object: ...
+def TFE_Py_TapeSetPossibleGradientTypes(arg0: object) -> object: ...
+def TFE_Py_TapeSetRecordOperation(arg0: object, arg1: object, arg2: object, arg3: object, arg4: object) -> object: ...
+def TFE_Py_TapeSetRecordOperationBackprop(arg0: object, arg1: object, arg2: object, arg3: object) -> object: ...
+def TFE_Py_TapeSetRecordOperationForwardprop(arg0: object, arg1: object, arg2: object, arg3: object, arg4: object) -> object: ...
+def TFE_Py_TapeSetRemove(arg0: object) -> None: ...
+def TFE_Py_TapeSetRestartOnThread() -> None: ...
+def TFE_Py_TapeSetShouldRecordBackprop(arg0: object) -> object: ...
+def TFE_Py_TapeSetStopOnThread() -> None: ...
+def TFE_Py_TapeVariableAccessed(arg0: object) -> None: ...
+def TFE_Py_TapeWatch(arg0: object, arg1: object) -> None: ...
+def TFE_Py_TapeWatchVariable(arg0: object, arg1: object) -> None: ...
+def TFE_Py_TapeWatchedVariables(arg0: object) -> object: ...
+def TFE_Py_TensorShapeOnDevice(arg0: object, arg1: int) -> object: ...
+def TFE_Py_TensorShapeSlice(arg0: object, arg1: int) -> object: ...
+def TFE_Py_UID() -> object: ...
+def TFE_Py_VariableWatcherNew() -> object: ...
+def TFE_Py_VariableWatcherRemove(arg0: object) -> None: ...
+def TFE_Py_VariableWatcherVariableAccessed(arg0: object) -> None: ...
+def TFE_Py_VariableWatcherWatchedVariables(arg0: object) -> object: ...
+def TFE_ReportErrorToCluster(arg0: object, arg1: int, arg2: str) -> None: ...
+def TFE_ResetMemoryStats(arg0: object, arg1: str) -> None: ...
+def TFE_SetLogicalCpuDevices(arg0: object, arg1: int, arg2: str) -> None: ...
+def TFE_ToDlpackCapsule(arg0: object): ...
+def TFE_WaitAtBarrier(arg0: object, arg1: str, arg2: int) -> None: ...
+def TF_DeleteDeviceList(arg0: TF_DeviceList) -> None: ...
+def TF_DeviceListCount(arg0: TF_DeviceList) -> int: ...
+def TF_DeviceListName(arg0: TF_DeviceList, arg1: int) -> str: ...
+def TF_DeviceListType(arg0: TF_DeviceList, arg1: int) -> str: ...
+def TF_EnableMlirBridge(arg0: bool) -> None: ...
+def TF_EnableXlaDevices() -> None: ...
+def TF_GetCompilerIr(arg0: object, arg1: str, arg2: str, arg3: str, arg4: object, arg5: object, arg6: str) -> bytes: ...
+def TF_GetDeviceDetails(arg0: int) -> dict[str, str]: ...
+def TF_GetXlaConstantFoldingDisabled() -> int: ...
+def TF_IsMlirBridgeEnabled() -> int: ...
+def TF_ListPhysicalDevices() -> object: ...
+def TF_ListPluggablePhysicalDevices() -> object: ...
+def TF_NewBufferFromString(arg0, arg1: int) -> TF_Buffer: ...
+def TF_PickUnusedPortOrDie() -> int: ...
+def TF_ResetJitCompilerFlags() -> None: ...
+def TF_SetTfXlaCpuGlobalJit(arg0: int) -> int: ...
+def TF_SetXlaAutoJitMode(arg0: str) -> None: ...
+def TF_SetXlaConstantFoldingDisabled(arg0: int) -> None: ...
+def TF_SetXlaEnableLazyCompilation(arg0: int) -> int: ...
+def TF_SetXlaMinClusterSize(arg0: int) -> None: ...

tensorflow/python/_pywrap_tfe_monitoring_reader.pyi

@@ -0,0 +1,21 @@
+# 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.
+# ==============================================================================
+
+class TFE_MonitoringCounterReader:
+    def __init__(self, *args, **kwargs) -> None: ...
+
+def TFE_MonitoringNewCounterReader(arg0: str) -> TFE_MonitoringCounterReader: ...
+def TFE_MonitoringReadCounter0(arg0: TFE_MonitoringCounterReader) -> int: ...
+def TFE_MonitoringReadCounter1(arg0: TFE_MonitoringCounterReader, arg1: str) -> int: ...

tensorflow/python/_pywrap_toco_api.pyi

@@ -0,0 +1,16 @@
+# 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.
+# ==============================================================================
+
+def TocoConvert(model_flags_proto_txt_raw: object, toco_flags_proto_txt_raw: object, input_contents_txt_raw: object, extended_return: bool = ...) -> object: ...

tensorflow/python/autograph/BUILD

@@ -0,0 +1,26 @@
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "autograph",
+    srcs = [
+        "__init__.py",
+    ],
+    strict_deps = True,
+    visibility = ["//visibility:public"],
+    deps = [
+        "//tensorflow/python/autograph/core:converter",
+        "//tensorflow/python/autograph/impl:api",
+        "//tensorflow/python/autograph/lang:directives",
+        "//tensorflow/python/autograph/lang:special_functions",
+        "//tensorflow/python/autograph/operators:__init__",
+        "//tensorflow/python/autograph/pyct/common_transformers",
+        "//tensorflow/python/autograph/utils:__init__",
+        "//tensorflow/python/autograph/utils:ag_logging",
+        "//tensorflow/python/util:all_util",
+    ],
+)

tensorflow/python/autograph/CONTRIBUTING.md

@@ -0,0 +1,105 @@
+# How to contribute
+
+We'd love to have your patches and contributions! Here are some guidelines. In general, we follow the [TensorFlow contributing guidelines](../../CONTRIBUTING.md), but have some [AutoGraph-specific style guidelines](STYLE_GUIDE.md). More details below.
+
+### Note to active contributors
+
+In preparation for TF 2.0, we moved the code base of AutoGraph from
+`tensorflow/contrib/autograph` to `tensorflow/python/autograph`. The move
+does not impact functionality, and AutoGraph will remain accessible under
+`tensorflow.contrib.autograph` until `tensorflow.contrib` is retired.
+
+## TensorFlow Code of Conduct
+Please review and follow the [TensorFlow Code of Conduct](../../CODE_OF_CONDUCT.md).
+
+## Contributor License Agreement
+
+Contributions to this project must be accompanied by a Contributor License
+Agreement. You (or your employer) retain the copyright to your contribution;
+this simply gives us permission to use and redistribute your contributions as
+part of the project. Head over to <https://cla.developers.google.com/> to see
+your current agreements on file or to sign a new one.
+
+You generally only need to submit a CLA once, so if you've already submitted one
+(even if it was for a different project), you probably don't need to do it
+again.
+
+## Code reviews
+
+All submissions, including submissions by project members, require review. We
+use GitHub pull requests for this purpose. Consult [GitHub
+Help](https://help.github.com/articles/about-pull-requests/) for more
+information on using pull requests.
+
+After a pull request is approved, we merge it. Note our merging process differs
+from GitHub in that we pull and submit the change into an internal version
+control system. This system automatically pushes a git commit to the GitHub
+repository (with credit to the original author) and closes the pull request.
+
+## Style
+
+See the [AutoGraph style guide](STYLE_GUIDE.md).
+
+## Unit tests
+
+Please include unit tests when contributing new features
+([example here](converters/continue_statements_test.py)), as they help to prove
+that your code works correctly, and guard against future breaking changes to
+lower the maintenance cost.
+
+It's also helpful to check that any
+changes you propose do not break existing unit tests. You can run tests using the command,
+
+```shell
+bazel test --config=opt --copt=-O3 --copt=-march=native \
+  //tensorflow/contrib/autograph/...
+```
+
+from the root of the `tensorflow` repository. For more details see the [main TensorFlow Contributing File](../../CONTRIBUTING.md)
+
+## Developer info
+
+### Module structure
+
+The graph below describes the dependencies between AutoGraph modules (not to be mistaken with the directory structure for these modules, which is flat):
+
+```dot
+digraph d_modules {
+  autograph [style=filled];
+  converters;
+  core;
+  impl;
+  lang;
+  operators;
+
+  autograph -> impl
+  autograph -> lang
+
+  impl -> converters
+  impl -> core
+  impl -> operators
+
+  lang -> operators
+
+  converters -> core
+  converters -> lang
+}
+```
+
+`autograph` is the sole user-visible module.
+
+A short description of the modules:
+
+ * `autograph`: the main module imported by the user and by the generated code; only contains declarations
+ * `impl`: high level code and the implementation of the api frontend
+ * `core`: base classes for the AutoGraph source code transformation logic; see in particular `converter.py`
+ * `lang`: special user-visible functions that serve as extensions to the Python language
+ * `converters`: collection of source code transformation modules specialized for particular AutoGraph features
+ * `operators`: collection of operators that AutoGraph overloads; these correspond to Python operators as well as Python syntactic structures, like control flow
+
+There are two additional modules, `pyct` and `utils`. These are independent of AutoGraph:
+
+ * `pyct`: a general purpose Python source code transformation library
+ * `utils`: the kitchen sync; deprecated
+
+Note: we have a long term plan to factor out an implementation of `impl` and `converters` that is independent of autograph, into a general purpose Python operator overloading library.

tensorflow/python/autograph/STYLE_GUIDE.md

@@ -0,0 +1,85 @@
+# AutoGraph Style Guide
+
+This page contains style decisions that developers should follow when
+contributing code to AutoGraph.
+
+## TensorFlow Style
+
+Follow the [TensorFlow style
+guide](https://www.tensorflow.org/community/contribute/code_style), the [documentation
+guide](https://www.tensorflow.org/community/contribute/docs) and the
+[Google Python style guide](https://google.github.io/styleguide/pyguide.html).
+
+Naming conventions:
+
+1.  The name is TensorFlow, not Tensorflow.
+2.  The name is AutoGraph, not Autograph.
+
+## AutoGraph Style
+
+Below are AutoGraph-specific conventions. In the event of conflict, it
+supersedes all previous conventions.
+
+1. __Types in docstrings.__ Use [PEP 484][https://www.python.org/dev/peps/pep-0484/]
+    notation to describe the type for args, return values and attributes.
+
+    Example:
+
+    ```
+    Args:
+      foo: Dict[str, List[int]], a dictionary of sorts
+    ```
+
+2.  __Citations in Docstrings.__ Write a `#### References` subsection at the
+    bottom of any docstring with citations. Use ICLR’s bibliography style to
+    write references; for example, order entries by the first author's last
+    name. Add a link to the paper if the publication is open source (ideally,
+    arXiv).
+
+    Write in-paragraph citations in general, e.g., [(Tran and Blei, 2018)][1].
+    Write in-text citations when the citation is a noun, e.g., [Tran and Blei
+    (2018)][1]. Write citations with more than two authors using et al., e.g.,
+    [(Tran et al., 2018)][1]. Separate multiple citations with semicolon, e.g.,
+    ([Tran and Blei, 2018][1]; [Gelman and Rubin, 1992][2]).
+
+    Examples:
+
+    ```none
+    #### References
+
+    # technical report
+    [1]: Tony Finch. Incremental calculation of weighted mean and variance.
+         _Technical Report_, 2009.
+         http://people.ds.cam.ac.uk/fanf2/hermes/doc/antiforgery/stats.pdf
+
+    # journal
+    [2]: Andrew Gelman and Donald B. Rubin. Inference from Iterative Simulation
+         Using Multiple Sequences. _Statistical Science_, 7(4):457-472, 1992.
+
+    # arXiv preprint
+    # use "et al." for papers with too many authors to maintain
+    [3]: Aaron van den Oord et al. Parallel WaveNet: Fast High-Fidelity Speech
+         Synthesis. _arXiv preprint arXiv:1711.10433_, 2017.
+         https://arxiv.org/abs/1711.10433
+
+    # conference
+    [4]: Yeming Wen, Paul Vicol, Jimmy Ba, Dustin Tran, and Roger Grosse.
+         Flipout: Efficient Pseudo-Independent Weight Perturbations on
+         Mini-Batches. In _International Conference on Learning
+         Representations_, 2018.
+         https://arxiv.org/abs/1803.04386
+    ```
+
+3.  Avoid LaTeX in docstrings.
+
+    *   It is not rendered in many (if not most) editors and can be hard to read
+        for both LaTeX experts and non-experts.
+
+4. Write docstring and comment math using ASCII friendly notation; python using
+    operators. E.g., `x**2` better than `x^2`, `x[i, j]` better than `x_{i,j}`,
+    `sum{ f(x[i]) : i=1...n }` better than `\sum_{i=1}^n f(x_i)` `int{sin(x) dx:
+    x in [0, 2 pi]}` better than `\int_0^{2\pi} sin(x) dx`.
+
+    *   The more we stick to python style, the more someone can
+        copy/paste/execute.
+    *   Python style is usually easier to read as ASCII.

tensorflow/python/autograph/__init__.py

@@ -0,0 +1,57 @@
+# Copyright 2016 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.
+# ==============================================================================
+"""Conversion of eager-style Python into TensorFlow graph code.
+
+NOTE: In TensorFlow 2.0, AutoGraph is automatically applied when using
+`tf.function`. This module contains lower-level APIs for advanced use.
+
+AutoGraph transforms a subset of Python which operates on TensorFlow objects
+into equivalent TensorFlow graph code. When executing the graph, it has the same
+effect as if you ran the original code in eager mode.
+Python code which doesn't operate on TensorFlow objects remains functionally
+unchanged, but keep in mind that `tf.function` only executes such code at trace
+time, and generally will not be consistent with eager execution.
+
+For more information, see the
+[AutoGraph reference documentation](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/autograph/g3doc/reference/index.md),
+and the [tf.function guide](https://www.tensorflow.org/guide/function#autograph_transformations).
+"""
+
+from tensorflow.python.util.all_util import remove_undocumented
+
+# TODO(mdan): Revisit this list once we finalize the generated code mechanism.
+_allowed_symbols = [
+    # Main API
+    'AutoGraphError',
+    'ConversionOptions',
+    'Feature',
+    'StackTraceMapper',
+    'convert',
+    'converted_call',
+    'do_not_convert',
+    'to_code',
+    'to_graph',
+    # Overloaded operators
+    'operators',
+    # Python language "extensions"
+    'set_element_type',
+    'set_loop_options',
+    'stack',
+    'tensor_list',
+    # Utilities: to be removed
+    'utils',
+]
+
+remove_undocumented(__name__, _allowed_symbols)

tensorflow/python/autograph/converters/BUILD

@@ -0,0 +1,402 @@
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:tensorflow.bzl", "py_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "slices",
+    srcs = ["slices.py"],
+    strict_deps = True,
+    visibility = ["//tensorflow:__subpackages__"],
+    deps = [
+        "//tensorflow/python/autograph/core:converter",
+        "//tensorflow/python/autograph/lang:directives",
+        "//tensorflow/python/autograph/pyct:templates",
+        "@pypi//gast",
+    ],
+)
+
+py_library(
+    name = "continue_statements",
+    srcs = ["continue_statements.py"],
+    strict_deps = True,
+    visibility = ["//tensorflow:__subpackages__"],
+    deps = [
+        "//tensorflow/python/autograph/core:converter",
+        "//tensorflow/python/autograph/pyct:anno",
+        "//tensorflow/python/autograph/pyct:qual_names",
+        "//tensorflow/python/autograph/pyct:templates",
+        "//tensorflow/python/autograph/pyct/static_analysis:activity",
+        "//tensorflow/python/autograph/pyct/static_analysis:annos",
+    ],
+)
+
+py_library(
+    name = "logical_expressions",
+    srcs = ["logical_expressions.py"],
+    strict_deps = True,
+    visibility = ["//tensorflow:__subpackages__"],
+    deps = [
+        "//tensorflow/python/autograph/core:converter",
+        "//tensorflow/python/autograph/pyct:parser",
+        "//tensorflow/python/autograph/pyct:templates",
+        "@pypi//gast",
+    ],
+)
+
+py_library(
+    name = "break_statements",
+    srcs = ["break_statements.py"],
+    strict_deps = True,
+    visibility = ["//tensorflow:__subpackages__"],
+    deps = [
+        "//tensorflow/python/autograph/core:converter",
+        "//tensorflow/python/autograph/pyct:anno",
+        "//tensorflow/python/autograph/pyct:qual_names",
+        "//tensorflow/python/autograph/pyct:templates",
+        "//tensorflow/python/autograph/pyct/static_analysis:activity",
+        "//tensorflow/python/autograph/pyct/static_analysis:annos",
+    ],
+)
+
+py_library(
+    name = "conditional_expressions",
+    srcs = ["conditional_expressions.py"],
+    strict_deps = True,
+    visibility = ["//tensorflow:__subpackages__"],
+    deps = [
+        "//tensorflow/python/autograph/core:converter",
+        "//tensorflow/python/autograph/pyct:parser",
+        "//tensorflow/python/autograph/pyct:templates",
+        "@pypi//gast",
+    ],
+)
+
+py_library(
+    name = "call_trees",
+    srcs = ["call_trees.py"],
+    strict_deps = True,
+    visibility = ["//tensorflow:__subpackages__"],
+    deps = [
+        "//tensorflow/python/autograph/core:converter",
+        "//tensorflow/python/autograph/pyct:anno",
+        "//tensorflow/python/autograph/pyct:parser",
+        "//tensorflow/python/autograph/pyct:qual_names",
+        "//tensorflow/python/autograph/pyct:templates",
+        "//tensorflow/python/autograph/utils:ag_logging",
+        "@pypi//gast",
+    ],
+)
+
+py_library(
+    name = "directives",
+    srcs = ["directives.py"],
+    strict_deps = True,
+    visibility = ["//tensorflow:__subpackages__"],
+    deps = [
+        "//tensorflow/python/autograph/core:converter",
+        "//tensorflow/python/autograph/lang:directives",
+        "//tensorflow/python/autograph/pyct:anno",
+        "//tensorflow/python/util:tf_inspect",
+        "@pypi//gast",
+    ],
+)
+
+py_library(
+    name = "variables",
+    srcs = ["variables.py"],
+    strict_deps = True,
+    visibility = ["//tensorflow:__subpackages__"],
+    deps = [
+        "//tensorflow/python/autograph/core:converter",
+        "//tensorflow/python/autograph/pyct:anno",
+        "//tensorflow/python/autograph/pyct:templates",
+        "@pypi//gast",
+    ],
+)
+
+py_library(
+    name = "control_flow",
+    srcs = ["control_flow.py"],
+    strict_deps = True,
+    visibility = ["//tensorflow:__subpackages__"],
+    deps = [
+        "//tensorflow/python/autograph/core:converter",
+        "//tensorflow/python/autograph/lang:directives",
+        "//tensorflow/python/autograph/pyct:anno",
+        "//tensorflow/python/autograph/pyct:cfg",
+        "//tensorflow/python/autograph/pyct:origin_info",
+        "//tensorflow/python/autograph/pyct:parser",
+        "//tensorflow/python/autograph/pyct:qual_names",
+        "//tensorflow/python/autograph/pyct:templates",
+        "//tensorflow/python/autograph/pyct/static_analysis:activity",
+        "//tensorflow/python/autograph/pyct/static_analysis:annos",
+        "//tensorflow/python/autograph/pyct/static_analysis:liveness",
+        "//tensorflow/python/autograph/pyct/static_analysis:reaching_definitions",
+        "//tensorflow/python/autograph/pyct/static_analysis:reaching_fndefs",
+        "@pypi//gast",
+    ],
+)
+
+py_library(
+    name = "functions",
+    srcs = ["functions.py"],
+    strict_deps = True,
+    visibility = ["//tensorflow:__subpackages__"],
+    deps = [
+        "//tensorflow/python/autograph/core:converter",
+        "//tensorflow/python/autograph/pyct:anno",
+        "//tensorflow/python/autograph/pyct:parser",
+        "//tensorflow/python/autograph/pyct:qual_names",
+        "//tensorflow/python/autograph/pyct:templates",
+        "//tensorflow/python/autograph/pyct/static_analysis:activity",
+        "//tensorflow/python/autograph/pyct/static_analysis:annos",
+        "@pypi//gast",
+    ],
+)
+
+py_library(
+    name = "lists",
+    srcs = ["lists.py"],
+    strict_deps = True,
+    visibility = ["//tensorflow:__subpackages__"],
+    deps = [
+        "//tensorflow/python/autograph/core:converter",
+        "//tensorflow/python/autograph/lang:directives",
+        "//tensorflow/python/autograph/pyct:anno",
+        "//tensorflow/python/autograph/pyct:parser",
+        "//tensorflow/python/autograph/pyct:qual_names",
+        "//tensorflow/python/autograph/pyct:templates",
+        "//tensorflow/python/autograph/pyct/static_analysis:activity",
+        "//tensorflow/python/autograph/pyct/static_analysis:annos",
+        "@pypi//gast",
+    ],
+)
+
+py_library(
+    name = "return_statements",
+    srcs = ["return_statements.py"],
+    strict_deps = True,
+    visibility = ["//tensorflow:__subpackages__"],
+    deps = [
+        "//tensorflow/python/autograph/core:converter",
+        "//tensorflow/python/autograph/pyct:anno",
+        "//tensorflow/python/autograph/pyct:parser",
+        "//tensorflow/python/autograph/pyct:qual_names",
+        "//tensorflow/python/autograph/pyct:templates",
+        "//tensorflow/python/autograph/pyct/static_analysis:activity",
+        "//tensorflow/python/autograph/pyct/static_analysis:annos",
+        "@pypi//gast",
+    ],
+)
+
+py_library(
+    name = "asserts",
+    srcs = ["asserts.py"],
+    strict_deps = True,
+    visibility = ["//tensorflow:__subpackages__"],
+    deps = [
+        "//tensorflow/python/autograph/core:converter",
+        "//tensorflow/python/autograph/pyct:templates",
+        "@pypi//gast",
+    ],
+)
+
+py_test(
+    name = "asserts_test",
+    srcs = ["asserts_test.py"],
+    strict_deps = True,
+    deps = [
+        ":asserts",
+        ":functions",
+        ":return_statements",
+        #internal proto upb dep
+        "//tensorflow/python/autograph/core:test_lib",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_test(
+    name = "break_statements_test",
+    srcs = ["break_statements_test.py"],
+    strict_deps = True,
+    deps = [
+        ":break_statements",
+        #internal proto upb dep
+        "//tensorflow/python/autograph/core:test_lib",
+        "//tensorflow/python/autograph/pyct:anno",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_test(
+    name = "call_trees_test",
+    srcs = ["call_trees_test.py"],
+    strict_deps = True,
+    deps = [
+        ":call_trees",
+        ":functions",
+        #internal proto upb dep
+        "//tensorflow/python/autograph/core:test_lib",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_test(
+    name = "conditional_expressions_test",
+    srcs = ["conditional_expressions_test.py"],
+    strict_deps = True,
+    deps = [
+        ":conditional_expressions",
+        #internal proto upb dep
+        "//tensorflow/python/autograph/core:test_lib",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_test(
+    name = "continue_statements_test",
+    srcs = ["continue_statements_test.py"],
+    strict_deps = True,
+    deps = [
+        ":continue_statements",
+        #internal proto upb dep
+        "//tensorflow/python/autograph/core:test_lib",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_test(
+    name = "control_flow_test",
+    srcs = ["control_flow_test.py"],
+    strict_deps = True,
+    deps = [
+        ":break_statements",
+        ":continue_statements",
+        ":control_flow",
+        #internal proto upb dep
+        "//third_party/py/numpy",
+        "//tensorflow/python/autograph/core:test_lib",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:sparse_tensor",
+        "//tensorflow/python/framework:tensor_util",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/util:nest",
+    ],
+)
+
+py_test(
+    name = "directives_test",
+    srcs = ["directives_test.py"],
+    strict_deps = True,
+    deps = [
+        ":directives",
+        #internal proto upb dep
+        "//tensorflow/python/autograph/core:test_lib",
+        "//tensorflow/python/autograph/lang:directives",
+        "//tensorflow/python/autograph/pyct:anno",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_test(
+    name = "functions_test",
+    srcs = ["functions_test.py"],
+    strict_deps = True,
+    deps = [
+        ":functions",
+        ":return_statements",
+        #internal proto upb dep
+        "//tensorflow/python/autograph/core:ag_ctx",
+        "//tensorflow/python/autograph/core:converter",
+        "//tensorflow/python/autograph/core:test_lib",
+        "//tensorflow/python/autograph/impl:api",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_test(
+    name = "lists_test",
+    srcs = ["lists_test.py"],
+    strict_deps = True,
+    deps = [
+        ":directives",
+        ":lists",
+        #internal proto upb dep
+        "//tensorflow/python/autograph/core:test_lib",
+        "//tensorflow/python/autograph/lang:directives",
+        "//tensorflow/python/autograph/lang:special_functions",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:tensor",
+        "//tensorflow/python/ops:array_ops_stack",
+        "//tensorflow/python/ops:list_ops",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_test(
+    name = "logical_expressions_test",
+    srcs = ["logical_expressions_test.py"],
+    strict_deps = True,
+    deps = [
+        ":logical_expressions",
+        #internal proto upb dep
+        "//tensorflow/python/autograph/core:test_lib",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_test(
+    name = "return_statements_test",
+    srcs = ["return_statements_test.py"],
+    strict_deps = True,
+    deps = [
+        ":functions",
+        ":return_statements",
+        #internal proto upb dep
+        "//tensorflow/python/autograph/core:test_lib",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_test(
+    name = "slices_test",
+    srcs = ["slices_test.py"],
+    strict_deps = True,
+    deps = [
+        ":directives",
+        ":slices",
+        #internal proto upb dep
+        "//tensorflow/python/autograph/core:test_lib",
+        "//tensorflow/python/autograph/lang:directives",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/ops:list_ops",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_test(
+    name = "variables_test",
+    srcs = ["variables_test.py"],
+    strict_deps = True,
+    deps = [
+        ":variables",
+        #internal proto upb dep
+        "//tensorflow/python/autograph/core:test_lib",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)

tensorflow/python/autograph/converters/asserts.py

@@ -0,0 +1,48 @@
+# Copyright 2016 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.
+# ==============================================================================
+"""Converts assert statements to their corresponding TF calls."""
+
+import gast
+
+from tensorflow.python.autograph.core import converter
+from tensorflow.python.autograph.pyct import templates
+
+
+class AssertTransformer(converter.Base):
+  """Transforms Assert nodes to Call so they can be handled as functions."""
+
+  def visit_Assert(self, node):
+    self.generic_visit(node)
+
+    # Note: The lone tf.Assert call will be wrapped with control_dependencies
+    # by side_effect_guards.
+    template = """
+      ag__.assert_stmt(test, lambda: msg)
+    """
+
+    if node.msg is None:
+      return templates.replace(
+          template,
+          test=node.test,
+          msg=gast.Constant('Assertion error', kind=None))
+    elif isinstance(node.msg, gast.Constant):
+      return templates.replace(template, test=node.test, msg=node.msg)
+    else:
+      raise NotImplementedError('can only convert string messages for now.')
+
+
+def transform(node, ctx):
+  node = AssertTransformer(ctx).visit(node)
+  return node

tensorflow/python/autograph/converters/asserts_test.py

@@ -0,0 +1,42 @@
+# 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 asserts module."""
+
+from tensorflow.python.autograph.converters import asserts
+from tensorflow.python.autograph.converters import functions
+from tensorflow.python.autograph.converters import return_statements
+from tensorflow.python.autograph.core import converter_testing
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import errors_impl
+from tensorflow.python.platform import test
+
+
+class AssertsTest(converter_testing.TestCase):
+
+  def test_basic(self):
+
+    def f(a):
+      assert a, 'testmsg'
+      return a
+
+    tr = self.transform(f, (functions, asserts, return_statements))
+
+    op = tr(constant_op.constant(False))
+    with self.assertRaisesRegex(errors_impl.InvalidArgumentError, 'testmsg'):
+      self.evaluate(op)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/autograph/converters/break_statements.py

@@ -0,0 +1,185 @@
+# 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.
+# ==============================================================================
+"""Lowers break statements to conditionals."""
+
+from tensorflow.python.autograph.core import converter
+from tensorflow.python.autograph.pyct import anno
+from tensorflow.python.autograph.pyct import qual_names
+from tensorflow.python.autograph.pyct import templates
+from tensorflow.python.autograph.pyct.static_analysis import activity
+from tensorflow.python.autograph.pyct.static_analysis.annos import NodeAnno
+
+
+class _Break(object):
+
+  def __init__(self):
+    self.used = False
+    self.control_var_name = None
+
+  def __repr__(self):
+    return 'used: %s, var: %s' % (self.used, self.control_var_name)
+
+
+class BreakTransformer(converter.Base):
+  """Canonicalizes break statements into additional conditionals."""
+
+  def visit_Break(self, node):
+    self.state[_Break].used = True
+    var_name = self.state[_Break].control_var_name
+    # TODO(mdan): This will fail when expanded inside a top-level else block.
+    template = """
+      var_name = True
+      continue
+    """
+    return templates.replace(template, var_name=var_name)
+
+  def _guard_if_present(self, block, var_name):
+    """Prevents the block from executing if var_name is set."""
+    if not block:
+      return block
+
+    template = """
+        if not var_name:
+          block
+      """
+    node = templates.replace(
+        template,
+        var_name=var_name,
+        block=block)
+    return node
+
+  def _process_body(self, nodes, break_var):
+    self.state[_Break].enter()
+    self.state[_Break].control_var_name = break_var
+    nodes = self.visit_block(nodes)
+    break_used = self.state[_Break].used
+    self.state[_Break].exit()
+    return nodes, break_used
+
+  def visit_While(self, node):
+    original_node = node
+    scope = anno.getanno(node, NodeAnno.BODY_SCOPE)
+    break_var = self.ctx.namer.new_symbol('break_', scope.referenced)
+
+    node.test = self.visit(node.test)
+    node.body, break_used = self._process_body(node.body, break_var)
+    # A break in the else clause applies to the containing scope.
+    node.orelse = self.visit_block(node.orelse)
+
+    if not break_used:
+      template = """
+        while test:
+          body
+        orelse
+      """
+      node = templates.replace(
+          template, test=node.test, body=node.body, orelse=node.orelse)
+
+      new_while_node = node[0]
+      anno.copyanno(original_node, new_while_node, anno.Basic.DIRECTIVES)
+
+      return node
+
+    # Python's else clause only triggers if the loop exited cleanly (e.g.
+    # break did not trigger).
+    guarded_orelse = self._guard_if_present(node.orelse, break_var)
+
+    template = """
+      var_name = False
+      while not var_name and test:
+        body
+      orelse
+    """
+    node = templates.replace(
+        template,
+        var_name=break_var,
+        test=node.test,
+        body=node.body,
+        orelse=guarded_orelse)
+
+    new_while_node = node[1]
+    anno.copyanno(original_node, new_while_node, anno.Basic.DIRECTIVES)
+
+    return node
+
+  def visit_For(self, node):
+    original_node = node
+    scope = anno.getanno(node, NodeAnno.BODY_SCOPE)
+    break_var = self.ctx.namer.new_symbol('break_', scope.referenced)
+
+    node.target = self.visit(node.target)
+    node.iter = self.visit(node.iter)
+    node.body, break_used = self._process_body(node.body, break_var)
+    # A break in the else clause applies to the containing scope.
+    node.orelse = self.visit_block(node.orelse)
+
+    if not break_used:
+      template = """
+        for target in iter_:
+          body
+        orelse
+      """
+      node = templates.replace(
+          template,
+          iter_=node.iter,
+          target=node.target,
+          body=node.body,
+          orelse=node.orelse)
+
+      new_for_node = node[0]
+      anno.copyanno(original_node, new_for_node, anno.Basic.EXTRA_LOOP_TEST)
+      anno.copyanno(original_node, new_for_node, anno.Basic.DIRECTIVES)
+
+      return node
+
+    # Python's else clause only triggers if the loop exited cleanly (e.g.
+    # break did not trigger).
+    guarded_orelse = self._guard_if_present(node.orelse, break_var)
+    extra_test = templates.replace_as_expression(
+        'not var_name', var_name=break_var)
+
+    # The extra test is hidden in the AST, which will confuse the static
+    # analysis. To mitigate that, we insert a no-op statement that ensures
+    # the control variable is marked as used.
+    # TODO(mdan): Use a marker instead, e.g. ag__.condition_loop_on(var_name)
+    template = """
+      var_name = False
+      for target in iter_:
+        (var_name,)
+        body
+      orelse
+    """
+    node = templates.replace(
+        template,
+        var_name=break_var,
+        iter_=node.iter,
+        target=node.target,
+        body=node.body,
+        orelse=guarded_orelse)
+
+    new_for_node = node[1]
+    anno.setanno(new_for_node, anno.Basic.EXTRA_LOOP_TEST, extra_test)
+    anno.copyanno(original_node, new_for_node, anno.Basic.DIRECTIVES)
+
+    return node
+
+
+def transform(node, ctx):
+  node = qual_names.resolve(node)
+  node = activity.resolve(node, ctx, None)
+
+  transformer = BreakTransformer(ctx)
+  node = transformer.visit(node)
+  return node

tensorflow/python/autograph/converters/break_statements_test.py

@@ -0,0 +1,170 @@
+# 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 break_statements module."""
+
+from tensorflow.python.autograph.converters import break_statements
+from tensorflow.python.autograph.core import converter_testing
+from tensorflow.python.autograph.pyct import anno
+from tensorflow.python.platform import test
+
+
+class BreakCanonicalizationTest(converter_testing.TestCase):
+
+  def assertTransformedEquivalent(self, f, *inputs):
+    tr = self.transform(f, break_statements)
+    self.assertEqual(f(*inputs), tr(*inputs))
+
+  def test_while_loop(self):
+
+    def f(x):
+      v = []
+      while x > 0:
+        x -= 1
+        if x % 2 == 0:
+          break
+        v.append(x)
+      return v
+
+    self.assertTransformedEquivalent(f, 0)
+    self.assertTransformedEquivalent(f, 1)
+    self.assertTransformedEquivalent(f, 4)
+
+  def test_while_loop_preserves_directives(self):
+
+    def f(x):
+      while x > 0:
+        x -= 1
+        if x % 2 == 0:
+          break
+
+    _, node, ctx = self.transform(f, (), include_ast=True)
+    fake_annotation = object()
+    anno.setanno(node.body[0], anno.Basic.DIRECTIVES, fake_annotation)
+    node = break_statements.transform(node, ctx)
+
+    self.assertIs(
+        anno.getanno(node.body[1], anno.Basic.DIRECTIVES), fake_annotation)
+
+  def test_for_loop(self):
+
+    def f(a):
+      v = []
+      for x in a:
+        x -= 1
+        if x % 2 == 0:
+          break
+        v.append(x)
+      return v
+
+    tr = self.transform(f, break_statements)
+
+    self.assertEqual([3], tr([5, 4]))
+
+  def test_for_loop_preserves_directives(self):
+
+    def f(a):
+      for x in a:
+        if x % 2 == 0:
+          break
+
+    _, node, ctx = self.transform(f, (), include_ast=True)
+    fake_annotation = object()
+    anno.setanno(node.body[0], anno.Basic.DIRECTIVES, fake_annotation)
+    node = break_statements.transform(node, ctx)
+    self.assertIs(
+        anno.getanno(node.body[1], anno.Basic.DIRECTIVES), fake_annotation)
+
+  def test_nested(self):
+
+    def f(x):
+      v = []
+      u = []
+      w = []
+      while x > 0:
+        x -= 1
+        if x % 2 == 0:
+          if x % 3 != 0:
+            u.append(x)
+          else:
+            w.append(x)
+            break
+        v.append(x)
+      return v, u, w
+
+    self.assertTransformedEquivalent(f, 0)
+    self.assertTransformedEquivalent(f, 3)
+    self.assertTransformedEquivalent(f, 11)
+
+  def test_nested_loops(self):
+
+    def f(x):
+      v = []
+      u = []
+      while x > 0:
+        x -= 1
+        y = x
+        while y > 0:
+          y -= 1
+          if y % 2 == 0:
+            break
+          u.append(y)
+        if x == 0:
+          break
+        v.append(x)
+      return v, u
+
+    self.assertTransformedEquivalent(f, 0)
+    self.assertTransformedEquivalent(f, 2)
+    self.assertTransformedEquivalent(f, 3)
+    self.assertTransformedEquivalent(f, 5)
+
+  def test_loop_orelse(self):
+
+    def f(x):
+      v = []
+      u = []
+      while x > 0:
+        x -= 1
+        y = x
+        while y > 1:
+          break
+        else:
+          u.append(y)
+          break
+        v.append(x)
+      return v, u
+
+    self.assertTransformedEquivalent(f, 0)
+    self.assertTransformedEquivalent(f, 2)
+    self.assertTransformedEquivalent(f, 3)
+
+  def test_multiple_correlated_breaks_with_side_effects(self):
+    def f(cond1):
+      lst = []
+      while True:
+        if cond1:
+          lst.append(1)
+        else:
+          break
+        if lst[-1] > 0:  # lst always has an element here
+          break
+      return lst
+
+    self.assertTransformedEquivalent(f, True)
+    self.assertTransformedEquivalent(f, False)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/autograph/converters/call_trees.py

@@ -0,0 +1,221 @@
+# Copyright 2016 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.
+# ==============================================================================
+"""Handles function calls, by generating compiled function names and calls.
+
+Note: this transformer does not rename the top level object being converted;
+that is the caller's responsibility.
+
+Requires function_scopes.
+"""
+
+import gast
+
+from tensorflow.python.autograph.core import converter
+from tensorflow.python.autograph.pyct import anno
+from tensorflow.python.autograph.pyct import parser
+from tensorflow.python.autograph.pyct import qual_names
+from tensorflow.python.autograph.pyct import templates
+from tensorflow.python.autograph.utils import ag_logging
+
+
+# TODO(mdan): Rename to FunctionCallsTransformer.
+
+
+class _Function(object):
+
+  no_root = True
+
+  def __init__(self):
+    self.context_name = None
+
+
+set_trace_warned = False
+
+
+class _ArgTemplateBuilder(object):
+  """Constructs a tuple representing the positional arguments in a call.
+
+  Example (yes, it's legal Python 3):
+
+      f(*args1, b, *args2, c, d)  ->  args1 + (b,) + args2 + (c, d)
+  """
+
+  def __init__(self):
+    self._arg_accumulator = []
+    self._argspec = []
+    self._finalized = False
+
+  def _consume_args(self):
+    if self._arg_accumulator:
+      self._argspec.append(
+          gast.Tuple(elts=self._arg_accumulator, ctx=gast.Load()))
+      self._arg_accumulator = []
+
+  def add_arg(self, a):
+    self._arg_accumulator.append(a)
+
+  def add_stararg(self, a):
+    self._consume_args()
+    self._argspec.append(
+        gast.Call(
+            gast.Name(
+                'tuple', ctx=gast.Load(), annotation=None, type_comment=None),
+            args=[a],
+            keywords=()))
+
+  def finalize(self):
+    self._consume_args()
+    self._finalized = True
+
+  def to_ast(self):
+    assert self._finalized
+    if self._argspec:
+      result = self._argspec[0]
+      for i in range(1, len(self._argspec)):
+        result = gast.BinOp(result, gast.Add(), self._argspec[i])
+      return result
+    return gast.Tuple([], gast.Load())
+
+
+class CallTreeTransformer(converter.Base):
+  """Transforms the call tree by renaming transformed symbols."""
+
+  def visit_Lambda(self, node):
+    if not anno.hasanno(node, 'function_context_name'):
+      # Lambda functions created during the conversion process have no
+      # context manager.
+      return self.generic_visit(node)
+    with self.state[_Function] as fn_scope:
+      fn_scope.context_name = anno.getanno(node, 'function_context_name')
+      return self.generic_visit(node)
+
+  def visit_FunctionDef(self, node):
+    # Decorators and arg defaults are part of the outer scope.
+    node.decorator_list = self.visit_block(node.decorator_list)
+    node.args.defaults = self.visit_block(node.args.defaults)
+    for i, d in enumerate(node.args.kw_defaults):
+      if d is not None:
+        node.args.kw_defaults[i] = self.visit(d)
+    with self.state[_Function] as fn_scope:
+      # Note: if the conversion process ever creates helper functions, this
+      # assumption will no longer hold.
+      assert anno.hasanno(node, 'function_context_name'), (
+          'The function_scopes converter always creates a scope for functions.')
+      fn_scope.context_name = anno.getanno(node, 'function_context_name')
+      node.body = self.visit_block(node.body)
+      if node.returns:
+        node.returns = self.visit(node.returns)
+      return node
+
+  def visit_With(self, node):
+    # Context manager calls (in node.items) are not converted.
+    node.body = self.visit_block(node.body)
+    return node
+
+  def _args_to_tuple(self, node):
+    """Ties together all positional and *arg arguments in a single tuple."""
+    # TODO(mdan): We could rewrite this to just a call to tuple(). Maybe better?
+    # For example for
+    #   f(a, b, *args)
+    # instead of writing:
+    #   (a, b) + args
+    # just write this?
+    #   tuple(a, b, *args)
+    builder = _ArgTemplateBuilder()
+    for a in node.args:
+      if isinstance(a, gast.Starred):
+        builder.add_stararg(a.value)
+      else:
+        builder.add_arg(a)
+    builder.finalize()
+    return builder.to_ast()
+
+  def _kwargs_to_dict(self, node):
+    """Ties together all keyword and **kwarg arguments in a single dict."""
+    if node.keywords:
+      return gast.Call(
+          gast.Name(
+              'dict', ctx=gast.Load(), annotation=None, type_comment=None),
+          args=(),
+          keywords=node.keywords)
+    else:
+      return parser.parse_expression('None')
+
+  def visit_Call(self, node):
+    full_name = str(anno.getanno(node.func, anno.Basic.QN, default=''))
+    function_context_name = self.state[_Function].context_name
+    node = self.generic_visit(node)
+
+    # TODO(mdan): Refactor converted_call as a 'Call' operator.
+
+    # Calls to the internal 'ag__' module are never converted (though their
+    # arguments might be).
+    if full_name.startswith('ag__.'):
+      return node
+
+    # Calls to the function context manager (inserted by function_scopes) are
+    # also safe.
+    if full_name.startswith(function_context_name + '.'):
+      return node
+
+    # Calls to pdb.set_trace or ipdb.set_trace are never converted. We don't use
+    # the normal mechanisms to bypass these literals because they are sensitive
+    # to the frame they are being called from.
+    # TODO(mdan): Generalize this to a "static allowlist" config.
+    if full_name in ('pdb.set_trace', 'ipdb.set_trace', 'breakpoint'):
+      global set_trace_warned
+      if not set_trace_warned:
+        # TODO(mdan): Update and shorten once available on tensorflow.org.
+        ag_logging.warning(
+            'Detected `pdb.set_trace()` in user code. The code'
+            ' generated by AutoGraph is not optimized for step-by-step'
+            ' debugging. See https://github.com/tensorflow/tensorflow/'
+            'blob/master/tensorflow/python/autograph/g3doc/reference/'
+            'debugging.md.')
+        set_trace_warned = True
+      return node
+
+    if (full_name == 'print' and
+        not self.ctx.user.options.uses(converter.Feature.BUILTIN_FUNCTIONS)):
+      return node
+
+    template = """
+      ag__.converted_call(func, args, kwargs, function_ctx)
+    """
+    new_call = templates.replace_as_expression(
+        template,
+        func=node.func,
+        args=self._args_to_tuple(node),
+        kwargs=self._kwargs_to_dict(node),
+        function_ctx=function_context_name)
+
+    return new_call
+
+
+def transform(node, ctx):
+  """Transform function call to the compiled counterparts.
+
+  Args:
+    node: AST
+    ctx: EntityContext
+  Returns:
+    A tuple (node, new_names):
+        node: The transformed AST
+        new_names: set(string), containing any newly-generated names
+  """
+  node = qual_names.resolve(node)
+
+  node = CallTreeTransformer(ctx).visit(node)
+  return node

tensorflow/python/autograph/converters/call_trees_test.py

@@ -0,0 +1,262 @@
+# 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 call_trees module."""
+
+import types
+
+from tensorflow.python.autograph.converters import call_trees
+from tensorflow.python.autograph.converters import functions
+from tensorflow.python.autograph.core import converter_testing
+from tensorflow.python.platform import test
+
+
+class MockConvertedCall(object):
+
+  def __init__(self):
+    self.calls = []
+
+  def __call__(self, f, args, kwargs, caller_fn_scope=None, options=None):
+    del caller_fn_scope, options
+    self.calls.append((args, kwargs))
+    kwargs = kwargs or {}
+    return f(*args, **kwargs)
+
+
+class CallTreesTest(converter_testing.TestCase):
+
+  def _transform_with_mock(self, f):
+    mock = MockConvertedCall()
+    tr = self.transform(
+        f, (functions, call_trees),
+        ag_overrides={'converted_call': mock})
+    return tr, mock
+
+  def test_function_no_args(self):
+
+    def f(f):
+      return f() + 20
+
+    tr, mock = self._transform_with_mock(f)
+
+    self.assertEqual(tr(lambda: 1), 21)
+    self.assertListEqual(mock.calls, [((), None)])
+
+  def test_function_with_expression_in_argument(self):
+
+    def f(f, g):
+      return f(g() + 20) + 4000
+
+    tr, mock = self._transform_with_mock(f)
+
+    self.assertEqual(tr(lambda x: x + 300, lambda: 1), 4321)
+    self.assertListEqual(mock.calls, [
+        ((), None),
+        ((21,), None),
+    ])
+
+  def test_function_with_call_in_argument(self):
+
+    def f(f, g):
+      return f(g()) + 300
+
+    tr, mock = self._transform_with_mock(f)
+
+    self.assertEqual(tr(lambda x: x + 20, lambda: 1), 321)
+    self.assertListEqual(mock.calls, [
+        ((), None),
+        ((1,), None),
+    ])
+
+  def test_function_chaining(self):
+
+    def get_one():
+      return 1
+
+    def f():
+      return get_one().__add__(20)
+
+    tr, mock = self._transform_with_mock(f)
+
+    self.assertEqual(tr(), 21)
+    self.assertListEqual(mock.calls, [
+        ((), None),
+        ((20,), None),
+    ])
+
+  def test_function_with_single_arg(self):
+
+    def f(f, a):
+      return f(a) + 20
+
+    tr, mock = self._transform_with_mock(f)
+
+    self.assertEqual(tr(lambda a: a, 1), 21)
+    self.assertListEqual(mock.calls, [((1,), None)])
+
+  def test_function_with_args_only(self):
+
+    def f(f, a, b):
+      return f(a, b) + 300
+
+    tr, mock = self._transform_with_mock(f)
+
+    self.assertEqual(tr(lambda a, b: a + b, 1, 20), 321)
+    self.assertListEqual(mock.calls, [((1, 20), None)])
+
+  def test_function_with_kwarg(self):
+
+    def f(f, a, b):
+      return f(a, c=b) + 300
+
+    tr, mock = self._transform_with_mock(f)
+
+    self.assertEqual(tr(lambda a, c: a + c, 1, 20), 321)
+    self.assertListEqual(mock.calls, [((1,), {'c': 20})])
+
+  def test_function_with_kwargs_starargs(self):
+
+    def f(f, a, *args, **kwargs):
+      return f(a, *args, **kwargs) + 5
+
+    tr, mock = self._transform_with_mock(f)
+
+    self.assertEqual(
+        tr(lambda *args, **kwargs: 7, 1, *[2, 3], **{
+            'b': 4,
+            'c': 5
+        }), 12)
+    self.assertListEqual(mock.calls, [((1, 2, 3), {'b': 4, 'c': 5})])
+
+  def test_function_with_starargs_only(self):
+
+    def g(*args):
+      return sum(args)
+
+    def f():
+      args = [1, 20, 300]
+      return g(*args) + 4000
+
+    tr, mock = self._transform_with_mock(f)
+
+    self.assertEqual(tr(), 4321)
+    self.assertListEqual(mock.calls, [((1, 20, 300), None)])
+
+  def test_function_with_starargs_mixed(self):
+
+    def g(a, b, c, d):
+      return a * 1000 + b * 100 + c * 10 + d
+
+    def f():
+      args1 = (1,)
+      args2 = [3]
+      return g(*args1, 2, *args2, 4)
+
+    tr, mock = self._transform_with_mock(f)
+
+    self.assertEqual(tr(), 1234)
+    self.assertListEqual(mock.calls, [((1, 2, 3, 4), None)])
+
+  def test_function_with_kwargs_keywords(self):
+
+    def f(f, a, b, **kwargs):
+      return f(a, b=b, **kwargs) + 5
+
+    tr, mock = self._transform_with_mock(f)
+
+    self.assertEqual(
+        tr(lambda *args, **kwargs: 7, 1, 2, **{'c': 3}), 12)
+    self.assertListEqual(mock.calls, [((1,), {'b': 2, 'c': 3})])
+
+  def test_function_with_multiple_kwargs(self):
+
+    def f(f, a, b, c, kwargs1, kwargs2):
+      return f(a, b=b, **kwargs1, c=c, **kwargs2) + 5
+
+    tr, mock = self._transform_with_mock(f)
+
+    self.assertEqual(
+        tr(lambda *args, **kwargs: 7, 1, 2, 3, {'d': 4}, {'e': 5}), 12)
+    self.assertListEqual(mock.calls, [((1,), {
+        'b': 2,
+        'c': 3,
+        'd': 4,
+        'e': 5
+    })])
+
+  def test_function_with_call_in_lambda_argument(self):
+
+    def h(l, a):
+      return l(a) + 4000
+
+    def g(a, *args):
+      return a + sum(args)
+
+    def f(h, g, a, *args):
+      return h(lambda x: g(x, *args), a)
+
+    tr, _ = self._transform_with_mock(f)
+
+    self.assertEqual(tr(h, g, 1, *(20, 300)), 4321)
+
+  def test_debugger_set_trace(self):
+
+    tracking_list = []
+
+    pdb = types.ModuleType('fake_pdb')
+    pdb.set_trace = lambda: tracking_list.append(1)
+
+    def f():
+      return pdb.set_trace()
+
+    tr, _ = self._transform_with_mock(f)
+
+    tr()
+    self.assertListEqual(tracking_list, [1])
+
+  def test_class_method(self):
+
+    class TestClass(object):
+
+      def other_method(self, x):
+        return x + 20
+
+      def test_method(self, a):
+        return self.other_method(a) + 300
+
+    tc = TestClass()
+    tr, mock = self._transform_with_mock(TestClass.test_method)
+
+    self.assertEqual(321, tr(tc, 1))
+    self.assertListEqual(mock.calls, [((1,), None)])
+
+  def test_object_method(self):
+
+    class TestClass(object):
+
+      def other_method(self, x):
+        return x + 20
+
+      def test_method(self, a):
+        return self.other_method(a) + 300
+
+    tc = TestClass()
+    tr, mock = self._transform_with_mock(tc.test_method)
+
+    self.assertEqual(321, tr(tc, 1))
+    self.assertListEqual(mock.calls, [((1,), None)])
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/autograph/converters/conditional_expressions.py

@@ -0,0 +1,46 @@
+# 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.
+# ==============================================================================
+"""Converts the ternary conditional operator."""
+
+import gast
+
+from tensorflow.python.autograph.core import converter
+from tensorflow.python.autograph.pyct import parser
+from tensorflow.python.autograph.pyct import templates
+
+
+class ConditionalExpressionTransformer(converter.Base):
+  """Converts conditional expressions to functional form."""
+
+  def visit_IfExp(self, node):
+    template = '''
+        ag__.if_exp(
+            test,
+            lambda: true_expr,
+            lambda: false_expr,
+            expr_repr)
+    '''
+    expr_repr = parser.unparse(node.test, include_encoding_marker=False).strip()
+    return templates.replace_as_expression(
+        template,
+        test=node.test,
+        true_expr=node.body,
+        false_expr=node.orelse,
+        expr_repr=gast.Constant(expr_repr, kind=None))
+
+
+def transform(node, ctx):
+  node = ConditionalExpressionTransformer(ctx).visit(node)
+  return node

tensorflow/python/autograph/converters/conditional_expressions_test.py

@@ -0,0 +1,48 @@
+# 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 conditional_expressions module."""
+
+from tensorflow.python.autograph.converters import conditional_expressions
+from tensorflow.python.autograph.core import converter_testing
+from tensorflow.python.platform import test
+
+
+class ConditionalExpressionsTest(converter_testing.TestCase):
+
+  def assertTransformedEquivalent(self, f, *inputs):
+    tr = self.transform(f, conditional_expressions)
+    self.assertEqual(f(*inputs), tr(*inputs))
+
+  def test_basic(self):
+
+    def f(x):
+      return 1 if x else 0
+
+    self.assertTransformedEquivalent(f, 0)
+    self.assertTransformedEquivalent(f, 3)
+
+  def test_nested_orelse(self):
+
+    def f(x):
+      y = x * x if x > 0 else x if x else 1
+      return y
+
+    self.assertTransformedEquivalent(f, -2)
+    self.assertTransformedEquivalent(f, 0)
+    self.assertTransformedEquivalent(f, 2)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/autograph/converters/continue_statements.py

@@ -0,0 +1,164 @@
+# 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.
+# ==============================================================================
+"""Canonicalizes continue statements by de-sugaring into a control boolean."""
+
+from tensorflow.python.autograph.core import converter
+from tensorflow.python.autograph.pyct import anno
+from tensorflow.python.autograph.pyct import qual_names
+from tensorflow.python.autograph.pyct import templates
+from tensorflow.python.autograph.pyct.static_analysis import activity
+from tensorflow.python.autograph.pyct.static_analysis.annos import NodeAnno
+
+
+class _Continue(object):
+
+  def __init__(self):
+    self.used = False
+    self.control_var_name = None
+
+  def __repr__(self):
+    return '<_Continue(used: {}, var: {})>'.format(self.used,
+                                                   self.control_var_name)
+
+
+class _Block(object):
+  """Tracks information about lexical blocks as they are visited in the AST.
+
+  Mainly, this object tracks the creation of block guards that replace
+  `continue` statements (e.g. `if not continue_:`).
+
+  Attributes:
+    create_guard_current: bool, whether to create a guard for the current
+      statement.
+    create_guard_next: bool, whether to create a guard for the next
+      statement.
+    is_loop_type: bool, whether this block is the body of a loop.
+  """
+
+  def __init__(self):
+    self.is_loop_type = False
+    self.create_guard_current = False
+    self.create_guard_next = False
+
+
+class ContinueCanonicalizationTransformer(converter.Base):
+  """Canonicalizes continue statements into additional conditionals."""
+
+  def visit_Continue(self, node):
+    self.state[_Continue].used = True
+    for block in reversed(self.state[_Block].stack):
+      # See ContinueCanonicalizationTest.test_multiple_continues for an example
+      # it's necessary to create guards for all enclosing affected blocks, not
+      # just that of the current block.
+      block.create_guard_next = True
+      if block.is_loop_type:
+        # continue only affects the innermost loop
+        break
+    template = """
+      var_name = True
+    """
+    return templates.replace(
+        template, var_name=self.state[_Continue].control_var_name)
+
+  def _postprocess_statement(self, node):
+    if self.state[_Continue].used:
+      block = self.state[_Block]
+      should_wrap_current = block.create_guard_current
+      # After processing propagate whether to guard the next statement
+      block.create_guard_current = block.create_guard_next
+      block.create_guard_next = False
+      if should_wrap_current:
+        template = """
+          if not var_name:
+            original_node
+        """
+        cond, = templates.replace(
+            template,
+            var_name=self.state[_Continue].control_var_name,
+            original_node=node)
+        return cond, cond.body
+    return node, None
+
+  def _visit_loop_body(self, node, nodes):
+    self.state[_Continue].enter()
+    self.state[_Block].enter()
+    self.state[_Block].is_loop_type = True
+    scope = anno.getanno(node, NodeAnno.BODY_SCOPE)
+    continue_var = self.ctx.namer.new_symbol('continue_', scope.referenced)
+    self.state[_Continue].control_var_name = continue_var
+
+    nodes = self.visit_block(nodes, after_visit=self._postprocess_statement)
+
+    if self.state[_Continue].used:
+      template = """
+        var_name = False
+      """
+      control_var_init = templates.replace(template, var_name=continue_var)
+      nodes = control_var_init + nodes
+
+    self.state[_Block].exit()
+    self.state[_Continue].exit()
+    return nodes
+
+  def _visit_non_loop_body(self, nodes):
+    self.state[_Block].enter()
+    nodes = self.visit_block(nodes, after_visit=self._postprocess_statement)
+    self.state[_Block].exit()
+    return nodes
+
+  def visit_While(self, node):
+    node.test = self.visit(node.test)
+    node.body = self._visit_loop_body(node, node.body)
+    # A continue in the else clause applies to the containing scope.
+    node.orelse = self._visit_non_loop_body(node.orelse)
+    return node
+
+  def visit_For(self, node):
+    node.target = self.generic_visit(node.target)
+    node.iter = self.generic_visit(node.iter)
+    node.body = self._visit_loop_body(node, node.body)
+    # A continue in the else clause applies to the containing scope.
+    node.orelse = self._visit_non_loop_body(node.orelse)
+    return node
+
+  def visit_If(self, node):
+    node.body = self._visit_non_loop_body(node.body)
+    node.orelse = self._visit_non_loop_body(node.orelse)
+    return node
+
+  def visit_With(self, node):
+    node.items = self.visit_block(node.items)
+    node.body = self._visit_non_loop_body(node.body)
+    return node
+
+  def visit_Try(self, node):
+    node.body = self._visit_non_loop_body(node.body)
+    node.orelse = self._visit_non_loop_body(node.orelse)
+    # In Python 3.8 and later continue is allowed in finally blocks
+    node.finalbody = self._visit_non_loop_body(node.finalbody)
+    node.handlers = self.visit_block(node.handlers)
+    return node
+
+  def visit_ExceptHandler(self, node):
+    node.body = self._visit_non_loop_body(node.body)
+    return node
+
+
+def transform(node, ctx):
+  node = qual_names.resolve(node)
+  node = activity.resolve(node, ctx, None)
+
+  node = ContinueCanonicalizationTransformer(ctx).visit(node)
+  return node

tensorflow/python/autograph/converters/continue_statements_test.py

@@ -0,0 +1,223 @@
+# 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 continue_statements module."""
+
+from tensorflow.python.autograph.converters import continue_statements
+from tensorflow.python.autograph.core import converter_testing
+from tensorflow.python.framework import ops
+from tensorflow.python.platform import test
+
+
+class ContinueCanonicalizationTest(converter_testing.TestCase):
+
+  def assertTransformedEquivalent(self, f, *inputs):
+    tr = self.transform(f, continue_statements)
+    self.assertEqual(f(*inputs), tr(*inputs))
+
+  def test_basic(self):
+
+    def f(x):
+      v = []
+      while x > 0:
+        x -= 1
+        if x % 2 == 0:
+          continue
+        v.append(x)
+      return v
+
+    self.assertTransformedEquivalent(f, 0)
+    self.assertTransformedEquivalent(f, 1)
+    self.assertTransformedEquivalent(f, 3)
+    self.assertTransformedEquivalent(f, 4)
+
+  def test_multiple_continues(self):
+
+    def f(x):
+      v = []
+      while x > 0:
+        x -= 1
+        if x > 1:
+          continue
+        if x > 2:
+          continue
+        v.append(x)
+      return v
+
+    self.assertTransformedEquivalent(f, 0)
+    self.assertTransformedEquivalent(f, 1)
+    self.assertTransformedEquivalent(f, 3)
+    self.assertTransformedEquivalent(f, 4)
+
+  def test_multiple_continues_in_nested_scope(self):
+
+    def f(a):
+      v = []
+      for x in a:
+        x -= 1
+        if x > 100:
+          continue
+        try:
+          raise ValueError('intentional')
+        except ValueError:
+          continue
+        v.append(x)
+      return v
+
+    self.assertTransformedEquivalent(f, [])
+    self.assertTransformedEquivalent(f, [1])
+    self.assertTransformedEquivalent(f, [2])
+    self.assertTransformedEquivalent(f, [1, 2, 3])
+
+  def test_for_loop(self):
+
+    def f(a):
+      v = []
+      for x in a:
+        x -= 1
+        if x % 2 == 0:
+          continue
+        v.append(x)
+      return v
+
+    self.assertTransformedEquivalent(f, [])
+    self.assertTransformedEquivalent(f, [1])
+    self.assertTransformedEquivalent(f, [2])
+    self.assertTransformedEquivalent(f, [1, 2, 3])
+
+  def test_nested_with(self):
+
+    def f(x):
+      v = []
+      while x > 0:
+        x -= 1
+        with ops.name_scope(''):
+          if x % 2 == 0:
+            continue
+        v.append(x)
+      return v
+
+    self.assertTransformedEquivalent(f, 0)
+    self.assertTransformedEquivalent(f, 1)
+    self.assertTransformedEquivalent(f, 3)
+    self.assertTransformedEquivalent(f, 4)
+
+  def test_nested_multiple_withs(self):
+
+    def f(x):
+      v = []
+      while x > 0:
+        x -= 1
+        with ops.name_scope(''):
+          if x % 2 == 0:
+            continue
+        with ops.name_scope(''):
+          v.append(x)
+        v.append(x)
+      return v
+
+    self.assertTransformedEquivalent(f, 0)
+    self.assertTransformedEquivalent(f, 1)
+    self.assertTransformedEquivalent(f, 3)
+    self.assertTransformedEquivalent(f, 4)
+
+  def test_nested_multiple_withs_and_statements(self):
+
+    def f(x):
+      v = []
+      while x > 0:
+        x -= 1
+        with ops.name_scope(''):
+          if x % 2 == 0:
+            continue
+          v.append(x)
+        v.append(x)
+        with ops.name_scope(''):
+          v.append(x)
+        v.append(x)
+      return v
+
+    self.assertTransformedEquivalent(f, 0)
+    self.assertTransformedEquivalent(f, 1)
+    self.assertTransformedEquivalent(f, 3)
+    self.assertTransformedEquivalent(f, 4)
+
+  def test_nested_multiple_withs_and_nested_withs(self):
+
+    def f(x):
+      v = []
+      while x > 0:
+        x -= 1
+        with ops.name_scope(''):
+          if x % 2 == 0:
+            continue
+          with ops.name_scope(''):
+            v.append(x)
+        v.append(x)
+        with ops.name_scope(''):
+          v.append(x)
+        v.append(x)
+      return v
+
+    self.assertTransformedEquivalent(f, 0)
+    self.assertTransformedEquivalent(f, 1)
+    self.assertTransformedEquivalent(f, 3)
+    self.assertTransformedEquivalent(f, 4)
+
+  def test_nested(self):
+
+    def f(x):
+      v = []
+      u = []
+      w = []
+      while x > 0:
+        x -= 1
+        if x % 2 == 0:
+          if x % 3 != 0:
+            u.append(x)
+          else:
+            w.append(x)
+            continue
+        v.append(x)
+      return v, u, w
+
+    self.assertTransformedEquivalent(f, 0)
+    self.assertTransformedEquivalent(f, 1)
+    self.assertTransformedEquivalent(f, 3)
+    self.assertTransformedEquivalent(f, 4)
+
+  def test_multiple_guarded_continues_with_side_effects(self):
+
+    def f(x):
+      def track(u, x):
+        u.append(x)
+        return x
+
+      u = []
+      v = []
+      while x > 0:
+        x -= 1
+        if track(u, x) > 1:
+          continue
+        if track(u, x) > 2:
+          continue
+        v.append(x)
+      return u, v
+
+    self.assertTransformedEquivalent(f, 3)
+    self.assertTransformedEquivalent(f, 2)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/autograph/converters/control_flow.py

@@ -0,0 +1,413 @@
+# Copyright 2016 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.
+# ==============================================================================
+"""Handles control flow statements: while, for, if."""
+
+import gast
+
+from tensorflow.python.autograph.core import converter
+from tensorflow.python.autograph.lang import directives
+from tensorflow.python.autograph.pyct import anno
+from tensorflow.python.autograph.pyct import cfg
+from tensorflow.python.autograph.pyct import origin_info
+from tensorflow.python.autograph.pyct import parser
+from tensorflow.python.autograph.pyct import qual_names
+from tensorflow.python.autograph.pyct import templates
+from tensorflow.python.autograph.pyct.static_analysis import activity
+from tensorflow.python.autograph.pyct.static_analysis import annos
+from tensorflow.python.autograph.pyct.static_analysis import liveness
+from tensorflow.python.autograph.pyct.static_analysis import reaching_definitions
+from tensorflow.python.autograph.pyct.static_analysis import reaching_fndefs
+
+
+class _Function(object):
+
+  scope = None
+
+
+class ControlFlowTransformer(converter.Base):
+  """Transforms control flow structures like loops an conditionals."""
+
+  def visit_Lambda(self, node):
+    with self.state[_Function] as fn:
+      fn.scope = anno.getanno(node, anno.Static.SCOPE)
+      return self.generic_visit(node)
+
+  def visit_FunctionDef(self, node):
+    with self.state[_Function] as fn:
+      fn.scope = anno.getanno(node, annos.NodeAnno.BODY_SCOPE)
+      return self.generic_visit(node)
+
+  def _create_nonlocal_declarations(self, vars_):
+    vars_ = set(vars_)
+    results = []
+    global_vars = self.state[_Function].scope.globals & vars_
+
+    if global_vars:
+      results.append(gast.Global([str(v) for v in global_vars]))
+
+    nonlocal_vars = [
+        v for v in vars_ if not v.is_composite() and v not in global_vars]
+    if nonlocal_vars:
+      results.append(gast.Nonlocal([str(v) for v in nonlocal_vars]))
+
+    return results
+
+  def _create_state_functions(
+      self, block_vars, nonlocal_declarations, getter_name, setter_name):
+    if not block_vars:
+      template = """
+        def getter_name():
+          return ()
+        def setter_name(block_vars):
+          pass
+      """
+      return templates.replace(
+          template, getter_name=getter_name, setter_name=setter_name)
+
+    guarded_block_vars = []
+    for v in block_vars:
+      if v.is_simple():
+        guarded_block_vars.append(v)
+      else:
+        guarded_block_vars.append(
+            templates.replace_as_expression(
+                'ag__.ldu(lambda: var_, name)',
+                var_=v,
+                name=gast.Constant(str(v), kind=None)))
+
+    template = """
+      def getter_name():
+        return guarded_state_vars,
+      def setter_name(vars_):
+        nonlocal_declarations
+        state_vars, = vars_
+    """
+    return templates.replace(
+        template,
+        nonlocal_declarations=nonlocal_declarations,
+        getter_name=getter_name,
+        guarded_state_vars=guarded_block_vars,
+        setter_name=setter_name,
+        state_vars=tuple(block_vars))
+
+  def _create_loop_options(self, node):
+    if not anno.hasanno(node, anno.Basic.DIRECTIVES):
+      return gast.Dict([], [])
+
+    loop_directives = anno.getanno(node, anno.Basic.DIRECTIVES)
+    if directives.set_loop_options not in loop_directives:
+      return gast.Dict([], [])
+
+    opts_dict = loop_directives[directives.set_loop_options]
+    str_keys, values = zip(*opts_dict.items())
+    keys = [gast.Constant(s, kind=None) for s in str_keys]
+    values = list(values)  # ast and gast don't play well with tuples.
+    return gast.Dict(keys, values)
+
+  def _create_undefined_assigns(self, undefined_symbols):
+    assignments = []
+    for s in undefined_symbols:
+      template = '''
+        var = ag__.Undefined(symbol_name)
+      '''
+      assignments += templates.replace(
+          template,
+          var=s,
+          symbol_name=gast.Constant(s.ssf(), kind=None))
+    return assignments
+
+  def _get_block_basic_vars(self, modified, live_in, live_out):
+    nonlocals = self.state[_Function].scope.nonlocals
+    basic_scope_vars = []
+    for s in modified:
+      if s.is_composite():
+        # TODO(mdan): Raise an error when this happens for a TF scope.
+        continue
+      # Variables not live into or out of the scope are considered local to the
+      # scope.
+      if s in live_in or s in live_out or s in nonlocals:
+        basic_scope_vars.append(s)
+      continue
+    return frozenset(basic_scope_vars)
+
+  def _get_block_composite_vars(self, modified, live_in):
+    # The scope variables corresponding to composite symbols (e.g. `self.x`).
+    composite_scope_vars = []
+    for s in modified:
+      if not s.is_composite():
+        continue
+      # Mutations made to objects created inside the scope will appear as writes
+      # to composite symbols. Because these mutations appear as modifications
+      # made to composite symbols, we check whether the composite's parent is
+      # actually live into the scope.
+      # Example:
+      #   while cond:
+      #     x = Foo()
+      #     x.foo = 2 * x.foo  # x.foo is live into the scope, but x is not.
+      #
+      # Note that some parents might not be symbols - for example, in x['foo'],
+      # 'foo' is a parent, but it's a literal, not a symbol. We don't check the
+      # liveness of literals.
+      support_set_symbols = tuple(
+          sss for sss in s.support_set if sss.is_symbol())
+      if not all(sss in live_in for sss in support_set_symbols):
+        continue
+      composite_scope_vars.append(s)
+    return frozenset(composite_scope_vars)
+
+  def _get_block_vars(self, node, modified):
+    """Determines the variables affected inside a control flow statement."""
+    defined_in = anno.getanno(node, anno.Static.DEFINED_VARS_IN)
+    live_in = anno.getanno(node, anno.Static.LIVE_VARS_IN)
+    live_out = anno.getanno(node, anno.Static.LIVE_VARS_OUT)
+    fn_scope = self.state[_Function].scope
+
+    basic_scope_vars = self._get_block_basic_vars(
+        modified,
+        live_in,
+        live_out)
+    composite_scope_vars = self._get_block_composite_vars(modified, live_in)
+    scope_vars = tuple(basic_scope_vars | composite_scope_vars)
+
+    # Variables that are modified inside the scope, but not defined
+    # before entering it. Only simple variables must be defined. The
+    # composite ones will be implicitly checked at runtime.
+    possibly_undefined = (
+        modified - defined_in - fn_scope.globals - fn_scope.nonlocals)
+    undefined = tuple(v for v in possibly_undefined if not v.is_composite())
+
+    # Variables that are modified inside the scope, and depend on values outside
+    # it.
+    input_only = basic_scope_vars & live_in - live_out
+
+    # Place the outputs first, then sort lexicographically.
+    scope_vars = sorted(scope_vars, key=lambda v: (v in input_only, v))
+    nouts = len(scope_vars) - len(input_only)
+
+    return scope_vars, undefined, nouts
+
+  def visit_If(self, node):
+    node = self.generic_visit(node)
+    body_scope = anno.getanno(node, annos.NodeAnno.BODY_SCOPE)
+    orelse_scope = anno.getanno(node, annos.NodeAnno.ORELSE_SCOPE)
+
+    cond_vars, undefined, nouts = self._get_block_vars(
+        node, body_scope.bound | orelse_scope.bound)
+
+    undefined_assigns = self._create_undefined_assigns(undefined)
+
+    nonlocal_declarations = self._create_nonlocal_declarations(cond_vars)
+
+    reserved = body_scope.referenced | orelse_scope.referenced
+    state_getter_name = self.ctx.namer.new_symbol('get_state', reserved)
+    state_setter_name = self.ctx.namer.new_symbol('set_state', reserved)
+    state_functions = self._create_state_functions(
+        cond_vars, nonlocal_declarations, state_getter_name, state_setter_name)
+
+    orelse_body = node.orelse
+    if not orelse_body:
+      orelse_body = [gast.Pass()]
+
+    template = """
+      state_functions
+      def body_name():
+        nonlocal_declarations
+        body
+      def orelse_name():
+        nonlocal_declarations
+        orelse
+      undefined_assigns
+      ag__.if_stmt(
+        test,
+        body_name,
+        orelse_name,
+        state_getter_name,
+        state_setter_name,
+        (symbol_names,),
+        nouts)
+    """
+    new_nodes = templates.replace(
+        template,
+        body=node.body,
+        body_name=self.ctx.namer.new_symbol('if_body', reserved),
+        orelse=orelse_body,
+        orelse_name=self.ctx.namer.new_symbol('else_body', reserved),
+        nonlocal_declarations=nonlocal_declarations,
+        nouts=gast.Constant(nouts, kind=None),
+        state_functions=state_functions,
+        state_getter_name=state_getter_name,
+        state_setter_name=state_setter_name,
+        symbol_names=tuple(gast.Constant(str(s), kind=None) for s in cond_vars),
+        test=node.test,
+        undefined_assigns=undefined_assigns)
+    origin_info.copy_origin(node, new_nodes[-1])
+    return new_nodes
+
+  def visit_While(self, node):
+    node = self.generic_visit(node)
+    body_scope = anno.getanno(node, annos.NodeAnno.BODY_SCOPE)
+
+    loop_vars, undefined, _ = self._get_block_vars(node, body_scope.bound)
+
+    undefined_assigns = self._create_undefined_assigns(undefined)
+
+    nonlocal_declarations = self._create_nonlocal_declarations(loop_vars)
+
+    reserved = body_scope.referenced
+    state_getter_name = self.ctx.namer.new_symbol('get_state', reserved)
+    state_setter_name = self.ctx.namer.new_symbol('set_state', reserved)
+    state_functions = self._create_state_functions(
+        loop_vars, nonlocal_declarations, state_getter_name, state_setter_name)
+
+    opts = self._create_loop_options(node)
+
+    template = """
+      state_functions
+      def body_name():
+        nonlocal_declarations
+        body
+      def test_name():
+        return test
+      undefined_assigns
+      ag__.while_stmt(
+          test_name,
+          body_name,
+          state_getter_name,
+          state_setter_name,
+          (symbol_names,),
+          opts)
+    """
+    new_nodes = templates.replace(
+        template,
+        body=node.body,
+        body_name=self.ctx.namer.new_symbol('loop_body', reserved),
+        nonlocal_declarations=nonlocal_declarations,
+        opts=opts,
+        state_functions=state_functions,
+        state_getter_name=state_getter_name,
+        state_setter_name=state_setter_name,
+        symbol_names=tuple(gast.Constant(str(s), kind=None) for s in loop_vars),
+        test=node.test,
+        test_name=self.ctx.namer.new_symbol('loop_test', reserved),
+        undefined_assigns=undefined_assigns)
+    origin_info.copy_origin(node, new_nodes[-1])
+    return new_nodes
+
+  def visit_For(self, node):
+    node = self.generic_visit(node)
+    body_scope = anno.getanno(node, annos.NodeAnno.BODY_SCOPE)
+    iter_scope = anno.getanno(node, annos.NodeAnno.ITERATE_SCOPE)
+
+    loop_vars, undefined, _ = self._get_block_vars(
+        node, body_scope.bound | iter_scope.bound)
+
+    undefined_assigns = self._create_undefined_assigns(undefined)
+
+    nonlocal_declarations = self._create_nonlocal_declarations(loop_vars)
+
+    reserved = body_scope.referenced | iter_scope.referenced
+    state_getter_name = self.ctx.namer.new_symbol('get_state', reserved)
+    state_setter_name = self.ctx.namer.new_symbol('set_state', reserved)
+    state_functions = self._create_state_functions(
+        loop_vars, nonlocal_declarations, state_getter_name, state_setter_name)
+
+    opts = self._create_loop_options(node)
+    opts.keys.append(gast.Constant('iterate_names', kind=None))
+    opts.values.append(gast.Constant(
+        parser.unparse(node.target, include_encoding_marker=False), kind=None))
+
+    if anno.hasanno(node, anno.Basic.EXTRA_LOOP_TEST):
+      extra_test = anno.getanno(node, anno.Basic.EXTRA_LOOP_TEST)
+      extra_test_name = self.ctx.namer.new_symbol(
+          'extra_test', reserved)
+      template = """
+        def extra_test_name():
+          nonlocal_declarations
+          return extra_test_expr
+      """
+      extra_test_function = templates.replace(
+          template,
+          extra_test_expr=extra_test,
+          extra_test_name=extra_test_name,
+          loop_vars=loop_vars,
+          nonlocal_declarations=nonlocal_declarations)
+    else:
+      extra_test_name = parser.parse_expression('None')
+      extra_test_function = []
+
+    # iterate_arg_name holds a single arg with the iterates, which may be a
+    # tuple.
+    iterate_arg_name = self.ctx.namer.new_symbol('itr', reserved)
+    template = """
+      iterates = iterate_arg_name
+    """
+    iterate_expansion = templates.replace(
+        template, iterate_arg_name=iterate_arg_name, iterates=node.target)
+    origin_info.copy_origin(node, iterate_expansion)
+
+    template = """
+      state_functions
+      def body_name(iterate_arg_name):
+        nonlocal_declarations
+        iterate_expansion
+        body
+      extra_test_function
+      undefined_assigns
+      ag__.for_stmt(
+          iterated,
+          extra_test_name,
+          body_name,
+          state_getter_name,
+          state_setter_name,
+          (symbol_names,),
+          opts)
+    """
+    new_nodes = templates.replace(
+        template,
+        body=node.body,
+        body_name=self.ctx.namer.new_symbol('loop_body', reserved),
+        extra_test_function=extra_test_function,
+        extra_test_name=extra_test_name,
+        iterate_arg_name=iterate_arg_name,
+        iterate_expansion=iterate_expansion,
+        iterated=node.iter,
+        nonlocal_declarations=nonlocal_declarations,
+        opts=opts,
+        symbol_names=tuple(gast.Constant(str(s), kind=None) for s in loop_vars),
+        state_functions=state_functions,
+        state_getter_name=state_getter_name,
+        state_setter_name=state_setter_name,
+        undefined_assigns=undefined_assigns)
+    origin_info.copy_origin(node, new_nodes[-1])
+    return new_nodes
+
+
+class AnnotatedDef(reaching_definitions.Definition):
+
+  def __init__(self):
+    super(AnnotatedDef, self).__init__()
+    self.directives = {}
+
+
+def transform(node, ctx):
+  graphs = cfg.build(node)
+  node = qual_names.resolve(node)
+  node = activity.resolve(node, ctx, None)
+  node = reaching_definitions.resolve(node, ctx, graphs)
+  node = reaching_fndefs.resolve(node, ctx, graphs)
+  node = liveness.resolve(node, ctx, graphs)
+
+  node = ControlFlowTransformer(ctx).visit(node)
+  return node

tensorflow/python/autograph/converters/control_flow_test.py

@@ -0,0 +1,786 @@
+# 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 control_flow module."""
+
+import collections
+
+import numpy as np
+
+from tensorflow.python.autograph.converters import break_statements
+from tensorflow.python.autograph.converters import continue_statements
+from tensorflow.python.autograph.converters import control_flow
+from tensorflow.python.autograph.core import converter_testing
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework import tensor_util
+from tensorflow.python.platform import test
+from tensorflow.python.util import nest
+
+
+for_unaffected_global = None
+for_mixed_globals_nonglobals = None
+for_test_global_local = None
+
+
+class ControlFlowTestBase(converter_testing.TestCase):
+
+  def assertValuesEqual(self, actual, expected):
+    values = nest.map_structure(
+        lambda x: self.evaluate(x) if tensor_util.is_tf_type(x) else x,
+        actual)
+    self.assertAllEqual(values, expected)
+
+  def assertTransformedResult(self, f, inputs, expected):
+    if not isinstance(inputs, tuple):
+      inputs = (inputs,)
+    tr = self.transform(f, control_flow)
+    returns = tr(*inputs)
+    self.assertValuesEqual(returns, expected)
+
+
+class NestedControlFlowTest(ControlFlowTestBase):
+
+  def test_basic(self):
+
+    def f(n):
+      i = 0
+      j = 0
+      s = 0
+      while i < n:
+        while j < i:
+          j += 3
+        u = i + j  # 'u' is not defined within the inner loop
+        s += u
+        i += 1
+        j = 0
+      return s, i, j, n
+
+    self.assertTransformedResult(f, constant_op.constant(5),
+                                 (25, 5, 0, 5))
+
+  def test_mixed_globals_nonglobals(self):
+
+    def f(n):
+      global for_mixed_globals_nonglobals
+      i = 0
+      j = 0
+      for_mixed_globals_nonglobals = 0
+      while i < n:
+        while j < i:
+          j += 3
+        u = i + j  # 'u' is not defined within the inner loop
+        for_mixed_globals_nonglobals += u
+        i += 1
+        j = 0
+      return for_mixed_globals_nonglobals, i, j, n
+
+    self.assertTransformedResult(f, constant_op.constant(5),
+                                 (25, 5, 0, 5))
+
+  def test_composite_state_complex(self):
+
+    class TestClassX(object):
+
+      def __init__(self, x):
+        self.x = x
+
+    class TestClassY(object):
+
+      def __init__(self, y):
+        self.y = y
+
+    def f(n):
+      tc = TestClassX(TestClassY({'z': TestClassX(n)}))
+      if n > 0:
+        while n > 0:
+          if n < 2:
+            tc.x.y['z'].x += 1
+          n -= 1
+      return n, tc
+
+    tr = self.transform(f, control_flow)
+
+    n, tc = tr(constant_op.constant(5))
+    self.assertValuesEqual((n, tc.x.y['z'].x), (0, 6))
+
+
+class WhileStatementTest(ControlFlowTestBase):
+
+  def test_basic(self):
+
+    def f(n):
+      i = 0
+      s = 0
+      while i < n:
+        s += i
+        i += 1
+      return s, i, n
+
+    self.assertTransformedResult(f, constant_op.constant(5), (10, 5, 5))
+
+  def test_single_output(self):
+
+    def f(n):
+      while n > 0:
+        n -= 1
+      return n
+
+    self.assertTransformedResult(f, constant_op.constant(5), 0)
+
+  def test_composite_state_attr(self):
+
+    class TestClass(object):
+
+      def __init__(self):
+        self.x = constant_op.constant(3)
+
+    def f(n):
+      tc = TestClass()
+      while n > 0:
+        tc.x += 1
+        n -= 1
+      return n
+
+    self.assertTransformedResult(f, constant_op.constant(5), 0)
+
+  def test_composite_state_slice(self):
+
+    def f(n):
+      d = {'a': n}
+      k = 'a'
+      while n > 0:
+        d[k] += 1
+        n -= 1
+      return d[k], n
+
+    self.assertTransformedResult(f, constant_op.constant(5), (10, 0))
+
+  def test_composite_state_literal_slice(self):
+
+    def f(n):
+      d = {'a': n}
+      while n > 0:
+        d['a'] += 1
+        n -= 1
+      return d['a'], n
+
+    self.assertTransformedResult(f, constant_op.constant(5), (10, 0))
+
+  def test_composite_state_attr_initialized_in_loop(self):
+
+    class TestClass(object):
+      pass
+
+    def f(n, x):
+      tc = TestClass()
+      while n < 5:
+        if n == 0:
+          tc.subattr = x
+        else:
+          tc.subattr = tc.subattr + 1
+        n += 1
+      return tc.subattr
+
+    self.assertTransformedResult(f, (0, constant_op.constant(10)), 14)
+    tr = self.transform(f, control_flow)
+    with self.assertRaisesRegex(
+        ValueError, "'tc.subattr' must be defined before the loop"):
+      tr(constant_op.constant(0), 0)
+
+  def test_composite_state_slice_initialized_in_loop(self):
+
+    def f(n, x):
+      d = {}
+      k = 'subkey'
+      while n < 5:
+        if n == 0:
+          d[k] = x
+        else:
+          d[k] = d[k] + 1
+        n += 1
+      return d
+
+    self.assertTransformedResult(f, (0, constant_op.constant(10)),
+                                 {'subkey': 14})
+    tr = self.transform(f, control_flow)
+    with self.assertRaisesRegex(
+        ValueError, r"'d\[k\]' must be defined before the loop"):
+      tr(constant_op.constant(0), 0)
+
+  def test_composite_state_literal_slice_initialized_in_loop(self):
+
+    def f(n, x):
+      d = {}
+      while n < 5:
+        if n == 0:
+          d['subkey'] = x
+        else:
+          d['subkey'] = d['subkey'] + 1
+        n += 1
+      return d
+
+    self.assertTransformedResult(f, (0, constant_op.constant(10)),
+                                 {'subkey': 14})
+    tr = self.transform(f, control_flow)
+    with self.assertRaisesRegex(
+        ValueError, r"'d\['subkey'\]' must be defined before the loop"):
+      tr(constant_op.constant(0), 0)
+
+  def test_composite_state_slice_aliased_to_local(self):
+
+    def f(n, x):
+      d = {}
+      while n < 5:
+        k = 'subkey'
+        d[k] = x + 1
+        n += 1
+      return d
+
+    self.assertTransformedResult(f, (0, constant_op.constant(10)),
+                                 {'subkey': 11})
+    tr = self.transform(f, control_flow)
+    # TODO(b/136999953): Better error message.
+    # Note that this error happens at execution time.
+    with self.assertRaises(errors.InaccessibleTensorError):
+      graph_fn = def_function.function(tr, autograph=False)
+      self.evaluate(
+          graph_fn(constant_op.constant(0), constant_op.constant(5)))
+
+  def test_local_composite_attr(self):
+
+    class TestClass(object):
+
+      def __init__(self):
+        self.x = constant_op.constant(3)
+
+    def f(n):
+      while n > 0:
+        tc = TestClass()
+        tc.x = tc.x
+        n -= 1
+      return n
+
+    self.assertTransformedResult(f, constant_op.constant(5), 0)
+
+  def test_local_composite_slice(self):
+
+    def f(n):
+      while n > 0:
+        d = {'x': n}
+        k = 'x'
+        d[k] = d[k]
+        n -= 1
+      return n
+
+    self.assertTransformedResult(f, constant_op.constant(5), 0)
+
+  def test_local_composite_literal_slice(self):
+
+    def f(n):
+      while n > 0:
+        d = {'x': n}
+        d['x'] = d['x']
+        n -= 1
+      return n
+
+    self.assertTransformedResult(f, constant_op.constant(5), 0)
+
+  def test_non_tensor_state(self):
+
+    # This class is ok to be in a tf.while's state.
+    class TestClass(collections.namedtuple('TestClass', ('x'))):
+      pass
+
+    def f(n):
+      tc = TestClass([constant_op.constant(0)])
+      while n > 0:
+        tc = TestClass([constant_op.constant(3)])
+        tc.x[0] = tc.x[0] + 1
+        n -= 1
+      return tc.x[0]
+
+    self.assertTransformedResult(f, constant_op.constant(5), 4)
+
+  def test_non_tensor_state_illegal_type(self):
+
+    class TestClass(object):
+
+      def __init__(self):
+        self.x = [constant_op.constant(3)]
+
+    def f(n):
+      while n > 0:
+        tc = TestClass()
+        tc.x[0] = tc.x[0] + 1
+        n -= 1
+      return tc.x[0]
+
+    tr = self.transform(f, control_flow)
+
+    # The tested function would require `tc` to become part of the while loop
+    # state, but TensorFlow doesn't support classes at the moment.
+    with self.assertRaisesRegex(
+        ValueError, 'tc.*must be defined before the loop'):
+      tr(constant_op.constant(5))
+
+  def test_dispatches_by_cond_only(self):
+
+    class TensorIncompatibleNumeric(object):
+      """Works in arithmetic expression, but errors out with TF ops."""
+
+      def __init__(self, val):
+        self.val = val
+
+      def __add__(self, other):
+        return TensorIncompatibleNumeric(self.val + other)
+
+    def f(n, s):
+      while n > 0:
+        n -= 1
+        s += n
+      return s
+
+    self.assertTransformedResult(f, (constant_op.constant(5), 0), 10)
+    tr = self.transform(f, control_flow)
+    # n alone controls the staging. When the loop is not staged, Python
+    # knows how to add the two objects. But when staged, tf.while will
+    # not know how to deal with the TensorIncompatibleNumeric object.
+    self.assertEqual(tr(5, TensorIncompatibleNumeric(0)).val, 10)
+    with self.assertRaises(TypeError):
+      tr(constant_op.constant(5), TensorIncompatibleNumeric(0))
+
+
+class IfStatementTest(ControlFlowTestBase):
+
+  def test_basic(self):
+
+    def f(n):
+      a = 0
+      b = 0
+      if n > 0:
+        a = -n
+      else:
+        b = 2 * n
+      return a, b
+
+    self.assertTransformedResult(f, constant_op.constant(1), (-1, 0))
+    self.assertTransformedResult(f, constant_op.constant(-1), (0, -2))
+
+  def test_sparse_tensor(self):
+
+    def f(cond, a):
+      if cond:
+        a = -a
+      return a
+
+    st = sparse_tensor.SparseTensor(
+        indices=((0,),), values=(0,), dense_shape=(1,))
+    self.assertTransformedResult(f, (st, constant_op.constant(1)), -1)
+    self.assertTransformedResult(f, (None, constant_op.constant(1)), 1)
+
+  def test_complex_outputs(self):
+
+    class TestClass(object):
+
+      def __init__(self, a, b):
+        self.a = a
+        self.b = b
+
+    def f(n, obj):
+      obj.a = 0
+      obj.b = 0
+      if n > 0:
+        obj.a = -n
+      else:
+        obj.b = 2 * n
+      return obj
+
+    tr = self.transform(f, control_flow)
+
+    res_obj = tr(constant_op.constant(1), TestClass(0, 0))
+    self.assertValuesEqual((res_obj.a, res_obj.b), (-1, 0))
+    res_obj = tr(constant_op.constant(-1), TestClass(0, 0))
+    self.assertValuesEqual((res_obj.a, res_obj.b), (0, -2))
+
+  def test_single_output(self):
+
+    def f(n):
+      if n > 0:
+        n = -n
+      return n
+
+    self.assertTransformedResult(f, constant_op.constant(1), -1)
+
+  def test_unbalanced(self):
+
+    def f(n):
+      if n > 0:
+        n = 3
+      return n
+
+    self.assertTransformedResult(f, constant_op.constant(2), 3)
+    self.assertTransformedResult(f, constant_op.constant(-3), -3)
+
+  def test_unbalanced_raising(self):
+
+    def f(n):
+      if n > 0:
+        n = n + 1
+        raise ValueError()
+      return n
+
+    self.assertTransformedResult(f, -3, -3)
+
+    tr = self.transform(f, control_flow)
+
+    with self.assertRaises(ValueError):
+      tr(1)
+
+  def test_local_var(self):
+
+    def f(n):
+      if n > 0:
+        b = 4
+        n = b + 1
+      return n
+
+    self.assertTransformedResult(f, constant_op.constant(1), 5)
+    self.assertTransformedResult(f, constant_op.constant(-1), -1)
+
+  def test_local_remains_local(self):
+
+    def f(n):
+      if n > 0:
+        b = 4
+        n = b + 1
+      return n
+
+    self.assertTransformedResult(f, constant_op.constant(1), 5)
+    self.assertTransformedResult(f, constant_op.constant(-1), -1)
+
+  def test_global_local(self):
+
+    def f(n):
+      if n > 0:
+        global for_test_global_local
+        if for_test_global_local is None:
+          for_test_global_local = 1
+        else:
+          for_test_global_local += 1
+        n += for_test_global_local
+      return n
+
+    tr = self.transform(f, control_flow)
+    assert for_test_global_local is None
+    self.assertEqual(tr(1), 2)
+    self.assertEqual(for_test_global_local, 1)
+
+  def test_no_outputs(self):
+
+    def f(n):
+      if n > 0:
+        b = 4  # pylint:disable=unused-variable
+      return n
+
+    self.assertTransformedResult(f, constant_op.constant(1), 1)
+    self.assertTransformedResult(f, constant_op.constant(-1), -1)
+
+  def test_created_outputs(self):
+
+    def f(i):
+      if i == 0:
+        result = i - 1
+      else:
+        result = i + 1
+      return result
+
+    self.assertTransformedResult(f, 0, -1)
+    self.assertTransformedResult(f, 1, 2)
+
+  def test_created_loop_local_outputs(self):
+
+    def f(n, x):
+      for i in n:
+        if i == 0:
+          result = i - 1
+        else:
+          result = i + 1
+        if result > 0:
+          x += 1
+      return x
+
+    self.assertTransformedResult(f, (range(5), 10), 14)
+
+  def test_created_loop_variable(self):
+
+    def f(n, x):
+      for i in n:
+        if i == 0:
+          result = i - 1
+        if i > 0:  # Using the result from previous iteration.
+          if result < 0:
+            x += 1
+      return x
+
+    self.assertTransformedResult(f, (range(5), 10), 14)
+
+  def test_unaffected_global(self):
+
+    global for_unaffected_global
+    for_unaffected_global = 3
+
+    def f(i):
+      global for_unaffected_global
+      if i == 0:
+        for_unaffected_global = i - 1
+      return for_unaffected_global
+
+    self.assertTransformedResult(f, 1, 3)
+    self.assertTransformedResult(f, 0, -1)
+    self.assertEqual(for_unaffected_global, -1)
+
+  def test_unaffected_nonlocal(self):
+
+    def f(i):
+      def inner_fn():
+        nonlocal n
+        if i == 0:
+          n = i - 1
+
+      n = 3
+      inner_fn()
+      return n
+
+    self.assertTransformedResult(f, 1, 3)
+    self.assertTransformedResult(f, 0, -1)
+
+  def test_output_defined_in_prior_except(self):
+
+    def f(i):
+      try:
+        raise ValueError()
+      except ValueError:
+        x = 1
+      if i == 0:
+        x = i - 1
+      return x
+
+    self.assertTransformedResult(f, 1, 1)
+    self.assertTransformedResult(f, 0, -1)
+
+  def test_unbalanced_multiple_composites(self):
+
+    class Foo(object):
+
+      def __init__(self):
+        self.b = 2
+        self.c = 3
+
+    def f(x, condition):
+
+      z = 5
+      if condition:
+        x.b = 7
+        x.c = 11
+        z = 13
+
+      return x.b, x.c, z
+
+    self.assertTransformedResult(f, (Foo(), constant_op.constant(True)),
+                                 (7, 11, 13))
+    self.assertTransformedResult(f, (Foo(), constant_op.constant(False)),
+                                 (2, 3, 5))
+
+  def test_unbalanced_composite(self):
+
+    class Foo(object):
+
+      def __init__(self):
+        self.b = 2
+
+    def f(x, condition):
+
+      z = 5
+      if condition:
+        x.b = 7
+        z = 13
+
+      return x.b, z
+
+    self.assertTransformedResult(f, (Foo(), constant_op.constant(True)),
+                                 (7, 13))
+    self.assertTransformedResult(f, (Foo(), constant_op.constant(False)),
+                                 (2, 5))
+
+
+class ForStatementTest(ControlFlowTestBase):
+
+  def test_basic(self):
+
+    def f(l):
+      s1 = 0
+      s2 = 0
+      for e in l:
+        s1 += e
+        s2 += e * e
+      return s1, s2
+
+    self.assertTransformedResult(f, constant_op.constant([1, 3]), (4, 10))
+    empty_vector = constant_op.constant([], shape=(0,), dtype=dtypes.int32)
+    self.assertTransformedResult(f, empty_vector, (0, 0))
+
+  def test_single_output(self):
+
+    def f(l):
+      s = 0
+      for e in l:
+        s += e
+      return s
+
+    self.assertTransformedResult(f, constant_op.constant([1, 3]), 4)
+    empty_vector = constant_op.constant([], shape=(0,), dtype=dtypes.int32)
+    self.assertTransformedResult(f, empty_vector, 0)
+
+  def test_iterated_expression(self):
+
+    eval_count = [0]
+
+    def count_evals(x):
+      eval_count[0] += 1
+      return x
+
+    def f(n):
+      s = 0
+      for e in count_evals(range(n)):
+        s += e
+      return s
+
+    tr = self.transform(f, control_flow)
+
+    self.assertEqual(tr(5), 10)
+    self.assertEqual(eval_count[0], 1)
+
+  def test_composite_state_initialized_in_loop(self):
+
+    class TestClass(object):
+      pass
+
+    def f(n, x):
+      tc = TestClass()
+      for i in n:
+        if i == 0:
+          tc.x = x
+        else:
+          tc.x = tc.x + i
+      return tc.x
+
+    self.assertTransformedResult(f, (range(5), constant_op.constant(10)), 20)
+    tr = self.transform(f, control_flow)
+
+    with self.assertRaisesRegex(
+        ValueError, "'tc.x' must be defined before the loop"):
+      tr(constant_op.constant(list(range(5))), 0)
+
+  def test_tuple_unpacking(self):
+
+    def f(x_list):
+      z = constant_op.constant(0)  # pylint:disable=undefined-variable
+      for i, x in enumerate(x_list):
+        z = z + x + i
+      return z
+
+    self.assertTransformedResult(f, [3, 3], 7)
+
+  def test_with_comprehension_in_body(self):
+
+    def f(l, n):
+      s = constant_op.constant(list(range(n)))
+      for _ in l:
+        s += constant_op.constant([a for a in range(n)])
+      return s
+
+    self.assertTransformedResult(f, (constant_op.constant([1, 2, 3]), 5),
+                                 np.array(range(5)) * 4)
+
+
+class AdvancedControlFlowTest(ControlFlowTestBase):
+
+  def assertTransformedEquivalent(self, f, *inputs):
+    tr = self.transform(
+        f, (break_statements, continue_statements, control_flow))
+    self.assertEqual(f(*inputs), tr(*inputs))
+
+  def test_while_with_else(self):
+
+    def f(x):
+      while x > 2:
+        x /= 2
+      else:
+        x += 1
+      return x
+
+    self.assertTransformedEquivalent(f, 4)
+    self.assertTransformedEquivalent(f, 2)
+
+  def test_while_with_else_and_break(self):
+
+    def f(cond1):
+      x = 8
+      while x > 2:
+        x /= 2
+        if cond1:
+          break
+      else:
+        x += 1
+      return x
+
+    self.assertTransformedEquivalent(f, True)
+    self.assertTransformedEquivalent(f, False)
+
+  def test_for_with_else(self):
+
+    def f(l):
+      res = 0
+      for x in l:
+        res += x
+      else:
+        res += 1
+      return res
+
+    self.assertTransformedEquivalent(f, [])
+    self.assertTransformedEquivalent(f, [1, 2])
+
+  def test_for_with_else_and_break(self):
+
+    def f(flag):
+      l = [1, 2, 3]
+      res = 0
+      for x in l:
+        res += x
+        if flag:
+          break
+      else:
+        res += 1
+      return res
+
+    self.assertTransformedEquivalent(f, True)
+    self.assertTransformedEquivalent(f, False)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/autograph/converters/directives.py

@@ -0,0 +1,177 @@
+# 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.
+# ==============================================================================
+"""Handles directives.
+
+This converter removes the directive functions from the code and moves the
+information they specify into AST annotations. It is a specialized form of
+static analysis, one that is specific to AutoGraph.
+
+Note that this requires that the actual directive functions are static - that
+is, they do not change at runtime. So if you do something like this:
+
+  tf.autograph.set_loop_options = <new function>
+
+Then the directive will may no longer be recognized. Furthermore, if the
+converted function is cached, such an action may be irreversible.
+"""
+
+import inspect
+
+import gast
+
+from tensorflow.python.autograph.core import converter
+from tensorflow.python.autograph.lang import directives
+from tensorflow.python.autograph.pyct import anno
+from tensorflow.python.util import tf_inspect
+
+
+STATIC_VALUE = 'static_value'
+"""Used for AST annotations, see visit_Name."""
+
+
+class _LoopScope(object):
+
+  def __init__(self):
+    self.ast_node = None
+    self.statements_visited = 0
+
+
+def _map_args(call_node, function):
+  """Maps AST call nodes to the actual function's arguments.
+
+  Args:
+    call_node: ast.Call
+    function: Callable[..., Any], the actual function matching call_node
+  Returns:
+    Dict[Text, ast.AST], mapping each of the function's argument names to
+    the respective AST node.
+  Raises:
+      ValueError: if the default arguments are not correctly set
+  """
+  args = call_node.args
+  kwds = {kwd.arg: kwd.value for kwd in call_node.keywords}
+  call_args = tf_inspect.getcallargs(function, *args, **kwds)
+
+  # Keyword arguments not specified in kwds will be mapped to their defaults,
+  # which are Python values. Since we don't currently have a way to transform
+  # those into AST references, we simply remove them. By convention, directives
+  # use UNSPECIFIED as default value for optional arguments. No other
+  # defaults should be present.
+  unexpected_defaults = []
+  for k in call_args:
+    if (k not in kwds
+        and call_args[k] not in args
+        and call_args[k] is not directives.UNSPECIFIED):
+      unexpected_defaults.append(k)
+  if unexpected_defaults:
+    raise ValueError('Unexpected keyword argument values, %s, for function %s'
+                     % (zip(unexpected_defaults,
+                            [call_args[k] for k in unexpected_defaults]),
+                        function))
+  return {k: v for k, v in call_args.items() if v is not directives.UNSPECIFIED}
+
+
+class DirectivesTransformer(converter.Base):
+  """Parses compiler directives and converts them into AST annotations."""
+
+  def _process_symbol_directive(self, call_node, directive):
+    if len(call_node.args) < 1:
+      raise ValueError('"%s" requires a positional first argument'
+                       ' as the target' % directive.__name__)
+    target = call_node.args[0]
+    defs = anno.getanno(target, anno.Static.ORIG_DEFINITIONS)
+    for def_ in defs:
+      def_.directives[directive] = _map_args(call_node, directive)
+    return call_node
+
+  def _process_statement_directive(self, call_node, directive):
+    if self.state[_LoopScope].statements_visited > 1:
+      raise ValueError(
+          '"%s" must be the first statement in the loop block' % (
+              directive.__name__))
+    if self.state[_LoopScope].level < 2:
+      raise ValueError(
+          '"%s" must be used inside a statement' % directive.__name__)
+    target = self.state[_LoopScope].ast_node
+    node_anno = anno.getanno(target, anno.Basic.DIRECTIVES, {})
+    node_anno[directive] = _map_args(call_node, directive)
+    anno.setanno(target, anno.Basic.DIRECTIVES, node_anno)
+    return call_node
+
+  def visit_Name(self, node):
+    node = self.generic_visit(node)
+    if isinstance(node.ctx, gast.Load):
+      defs = anno.getanno(node, anno.Static.DEFINITIONS, ())
+      is_defined = bool(defs)
+      if not is_defined and node.id in self.ctx.info.namespace:
+        anno.setanno(node, STATIC_VALUE, self.ctx.info.namespace[node.id])
+    return node
+
+  def visit_Attribute(self, node):
+    node = self.generic_visit(node)
+    parent_val = anno.getanno(node.value, STATIC_VALUE, default=None)
+    if parent_val is not None and inspect.ismodule(parent_val):
+      if hasattr(parent_val, node.attr):
+        anno.setanno(node, STATIC_VALUE, getattr(parent_val, node.attr))
+    return node
+
+  def visit_Assign(self, node):
+    self.state[_LoopScope].statements_visited += 1
+    return self.generic_visit(node)
+
+  def visit_AugAssign(self, node):
+    self.state[_LoopScope].statements_visited += 1
+    return self.generic_visit(node)
+
+  def visit_Expr(self, node):
+    self.state[_LoopScope].statements_visited += 1
+    node = self.generic_visit(node)
+    if isinstance(node.value, gast.Call):
+      call_node = node.value
+      static_val = anno.getanno(call_node.func, STATIC_VALUE, default=None)
+      if static_val is not None:
+        # Note: directive calls are not output in the generated code, hence
+        # the removal from the code by returning None.
+
+        if static_val is directives.set_element_type:
+          self._process_symbol_directive(call_node, static_val)
+          return None
+        elif static_val is directives.set_loop_options:
+          self._process_statement_directive(call_node, static_val)
+          return None
+    return node
+
+  # TODO(mdan): This will be insufficient for other control flow.
+  # That means that if we ever have a directive that affects things other than
+  # loops, we'll need support for parallel scopes, or have multiple converters.
+  def _track_and_visit_loop(self, node):
+    self.state[_LoopScope].enter()
+    self.state[_LoopScope].ast_node = node
+    node = self.generic_visit(node)
+    # Edge case: a loop with just one directive statement would become empty.
+    if not node.body:
+      node.body = [gast.Pass()]
+    self.state[_LoopScope].exit()
+    return node
+
+  def visit_While(self, node):
+    return self._track_and_visit_loop(node)
+
+  def visit_For(self, node):
+    return self._track_and_visit_loop(node)
+
+
+def transform(node, ctx):
+  return DirectivesTransformer(ctx).visit(node)

tensorflow/python/autograph/converters/directives_test.py

@@ -0,0 +1,133 @@
+# 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 directives module."""
+
+from tensorflow.python.autograph.converters import directives as directives_converter
+from tensorflow.python.autograph.core import converter_testing
+from tensorflow.python.autograph.lang import directives
+from tensorflow.python.autograph.pyct import anno
+from tensorflow.python.platform import test
+
+
+class DirectivesTest(converter_testing.TestCase):
+
+  def test_local_target(self):
+
+    def f():
+      l = []
+      string_var = 0
+      directives.set_element_type(l, 'a', string_var)
+
+    _, node, _ = self.transform(f, directives_converter, include_ast=True)
+
+    def_, = anno.getanno(node.body[0].targets[0],
+                         anno.Static.DEFINITIONS)
+    d = def_.directives[directives.set_element_type]
+    self.assertEqual(d['dtype'].value, 'a')
+    self.assertEqual(d['shape'].id, 'string_var')
+
+  def test_argument_target(self):
+
+    def f(a):
+      directives.set_element_type(a, 1, shape=2)
+      pass
+
+    _, node, _ = self.transform(f, directives_converter, include_ast=True)
+
+    def_, = anno.getanno(node.args.args[0], anno.Static.DEFINITIONS)
+    d = def_.directives[directives.set_element_type]
+    self.assertEqual(d['dtype'].value, 1)
+    self.assertEqual(d['shape'].value, 2)
+
+  def test_loop_target(self):
+
+    def f():
+      a = True
+      while True:
+        directives.set_loop_options(parallel_iterations=10, back_prop=a)  # pylint: disable=unexpected-keyword-arg
+        pass
+
+    _, node, _ = self.transform(f, directives_converter, include_ast=True)
+
+    d = anno.getanno(node.body[1], anno.Basic.DIRECTIVES)
+    d = d[directives.set_loop_options]
+    self.assertEqual(d['parallel_iterations'].value, 10)
+    self.assertEqual(d['back_prop'].id, 'a')
+    self.assertNotIn('swap_memory', d)
+
+  def test_loop_target_no_loop(self):
+
+    def f():
+      directives.set_loop_options()
+      pass
+
+    with self.assertRaisesRegex(ValueError, 'must be used inside a statement'):
+      self.transform(f, directives_converter, include_ast=True)
+
+  def test_loop_target_not_first(self):
+
+    def f():
+      a = 1
+      while True:
+        a = 2
+        directives.set_loop_options(parallel_iterations=10, back_prop=a)  # pylint: disable=unexpected-keyword-arg
+
+    with self.assertRaisesRegex(ValueError, 'must be the first statement'):
+      self.transform(f, directives_converter, include_ast=True)
+
+  def test_value_verification_does_not_trigger_properties(self):
+
+    self_test = self
+
+    class TestClass(object):
+
+      @property
+      def b(self):
+        self_test.fail('This should never be evaluated')
+
+    tc = TestClass()
+
+    def f():
+      return tc.b + 1
+
+    _, node, _ = self.transform(f, directives_converter, include_ast=True)
+
+    self.assertIsNotNone(node)
+
+  def test_value_verification_does_not_trigger_getattr(self):
+
+    class TestClass(object):
+
+      def __init__(self):
+        self.getattr_called = False
+
+      def __getattr__(self, _):
+        # Note: seems that any exception raised here is absorbed by hasattr.
+        # So we can't call test.fail or raise.
+        self.getattr_called = True
+
+    tc = TestClass()
+
+    def f():
+      return tc.b + 1
+
+    _, node, _ = self.transform(f, directives_converter, include_ast=True)
+
+    self.assertIsNotNone(node)
+    self.assertFalse(tc.getattr_called)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/autograph/converters/functions.py

@@ -0,0 +1,134 @@
+# 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.
+# ==============================================================================
+"""Converts function definitions and lambdas by adding necessary boilerplate."""
+
+import gast
+
+from tensorflow.python.autograph.core import converter
+from tensorflow.python.autograph.pyct import anno
+from tensorflow.python.autograph.pyct import parser
+from tensorflow.python.autograph.pyct import qual_names
+from tensorflow.python.autograph.pyct import templates
+from tensorflow.python.autograph.pyct.static_analysis import activity
+from tensorflow.python.autograph.pyct.static_analysis import annos
+
+
+class _Function(object):
+
+  def __init__(self):
+    self.context_name = None
+
+
+class FunctionTransformer(converter.Base):
+  """Wraps function bodies around autograph-specific boilerplate."""
+
+  def _function_scope_options(self, fn_scope):
+    """Returns the options with which to create function scopes."""
+    # Top-level function receive the options that were directly requested.
+    # All others receive the options corresponding to a recursive conversion.
+    # Note: this mainly controls the user_requested flag, which is important
+    # primarily because the FunctionScope context also creates a
+    # ControlStatusCtx(autograph=ENABLED) when user_requested is True. See
+    # function_wrappers.py.
+    if fn_scope.level == 2:
+      return self.ctx.user.options
+    return self.ctx.user.options.call_options()
+
+  def visit_Lambda(self, node):
+    with self.state[_Function] as fn_scope:
+      node = self.generic_visit(node)
+
+      # TODO(mdan): Fix the tests so that we can always add this decorator.
+      if fn_scope.level > 2:
+        return templates.replace_as_expression(
+            'ag__.autograph_artifact(l)', l=node)
+
+      scope = anno.getanno(node, anno.Static.SCOPE)
+      function_context_name = self.ctx.namer.new_symbol('lscope',
+                                                        scope.referenced)
+      fn_scope.context_name = function_context_name
+      anno.setanno(node, 'function_context_name', function_context_name)
+
+      template = """
+        ag__.with_function_scope(
+            lambda function_context: body, function_context_name, options)
+      """
+      node.body = templates.replace_as_expression(
+          template,
+          options=self._function_scope_options(fn_scope).to_ast(),
+          function_context=function_context_name,
+          function_context_name=gast.Constant(function_context_name, kind=None),
+          body=node.body)
+
+      return node
+
+  def visit_FunctionDef(self, node):
+    with self.state[_Function] as fn_scope:
+      scope = anno.getanno(node, annos.NodeAnno.BODY_SCOPE)
+
+      function_context_name = self.ctx.namer.new_symbol('fscope',
+                                                        scope.referenced)
+      fn_scope.context_name = function_context_name
+      anno.setanno(node, 'function_context_name', function_context_name)
+
+      node = self.generic_visit(node)
+
+      if fn_scope.level <= 2:
+        # Top-level functions lose their decorator because the conversion is
+        # always just-in-time and by the time it happens the decorators are
+        # already set to be applied.
+        node.decorator_list = []
+      else:
+        # TODO(mdan): Fix the tests so that we can always add this decorator.
+        # Inner functions are converted already, so we insert a decorator to
+        # prevent double conversion. Double conversion would work too, but this
+        # saves the overhead.
+        node.decorator_list.append(
+            parser.parse_expression('ag__.autograph_artifact'))
+
+      docstring_node = None
+      if node.body:
+        first_statement = node.body[0]
+        if (isinstance(first_statement, gast.Expr) and
+            isinstance(first_statement.value, gast.Constant)):
+          docstring_node = first_statement
+          node.body = node.body[1:]
+
+      template = """
+        with ag__.FunctionScope(
+            function_name, context_name, options) as function_context:
+          body
+      """
+      wrapped_body = templates.replace(
+          template,
+          function_name=gast.Constant(node.name, kind=None),
+          context_name=gast.Constant(function_context_name, kind=None),
+          options=self._function_scope_options(fn_scope).to_ast(),
+          function_context=function_context_name,
+          body=node.body)
+
+      if docstring_node is not None:
+        wrapped_body = [docstring_node] + wrapped_body
+
+      node.body = wrapped_body
+
+      return node
+
+
+def transform(node, ctx):
+  node = qual_names.resolve(node)
+  node = activity.resolve(node, ctx, None)
+
+  return FunctionTransformer(ctx).visit(node)

tensorflow/python/autograph/converters/functions_test.py

@@ -0,0 +1,127 @@
+# 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 functions module."""
+
+from tensorflow.python.autograph.converters import functions
+from tensorflow.python.autograph.converters import return_statements
+from tensorflow.python.autograph.core import ag_ctx
+from tensorflow.python.autograph.core import converter
+from tensorflow.python.autograph.core import converter_testing
+from tensorflow.python.autograph.impl import api
+from tensorflow.python.framework import constant_op
+from tensorflow.python.platform import test
+
+
+class FunctionTransformer(converter_testing.TestCase):
+
+  def test_basic(self):
+
+    def f(l):
+      """Docstring."""
+      a = 1
+      l += a
+      return l
+
+    tr = self.transform(f, functions)
+
+    result_op = tr(constant_op.constant(1))
+    self.assertIn('f/', result_op.op.name)
+    self.assertEqual('Docstring.', tr.__doc__)
+
+  def test_multiline_docstring(self):
+
+    def f():
+      """First sentence.
+
+      Second sentence.
+
+      Returns:
+        Something.
+      """
+      return constant_op.constant(1)
+
+    tr = self.transform(f, functions)
+
+    result_op = tr()
+    self.assertIn('f/', result_op.op.name)
+    self.assertIn('First sentence.', tr.__doc__)
+    self.assertIn('Second sentence.', tr.__doc__)
+
+  def test_nested_functions(self):
+
+    def f(l):
+
+      def inner_fn(i):
+        return i + 1
+
+      l += 1
+      return l, inner_fn(l)
+
+    tr = self.transform(f, (functions, return_statements))
+
+    first, second = tr(constant_op.constant(1))
+    self.assertIn('f/', first.op.name)
+    self.assertNotIn('inner_fn', first.op.name)
+    self.assertIn('f/inner_fn/', second.op.inputs[0].name)
+
+  def test_conversion_context_preserves_in_inner_functions(self):
+
+    def inner_fn_callee():
+      self.assertEqual(
+          ag_ctx.control_status_ctx().status, ag_ctx.Status.DISABLED)
+
+    def f():
+      def inner_fn():
+        inner_fn_callee()
+      with ag_ctx.ControlStatusCtx(
+          ag_ctx.Status.DISABLED, converter.ConversionOptions(recursive=True)):
+        inner_fn()
+
+    tr = self.transform(f, functions)
+
+    tr()
+
+  def test_method(self):
+
+    class TestClass(object):
+
+      def f(self, l):
+
+        def inner_fn(i):
+          return i + 1
+
+        l += 1
+        return l, inner_fn(l)
+
+    tr = self.transform(TestClass.f, (functions, return_statements))
+
+    first, second = tr(TestClass(), constant_op.constant(1))
+    self.assertIn('f/', first.op.name)
+    self.assertNotIn('inner_fn', first.op.name)
+    self.assertIn('f/inner_fn/', second.op.inputs[0].name)
+
+  def test_lambda_in_return_value(self):
+
+    def f():
+      return lambda x: x + 1
+
+    tr = self.transform(f, functions)
+
+    result_l = tr()
+    self.assertTrue(api.is_autograph_artifact(result_l))
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/autograph/converters/lists.py

@@ -0,0 +1,239 @@
+# Copyright 2016 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.
+# ==============================================================================
+"""Converter for list operations.
+
+This includes converting Python lists to TensorArray/TensorList.
+"""
+
+# TODO(mdan): Elaborate the logic here.
+# TODO(mdan): Does it even make sense to attempt to try to use TAs?
+# The current rule (always convert to TensorArray) is naive and insufficient.
+# In general, a better mechanism could look like:
+#   * convert to TensorList by default
+#   * leave as Python list if the user explicitly forbids it
+#   * convert to TensorArray only when complete write once behavior can be
+#     guaranteed (e.g. list comprehensions)
+
+import gast
+
+from tensorflow.python.autograph.core import converter
+from tensorflow.python.autograph.lang import directives
+from tensorflow.python.autograph.pyct import anno
+from tensorflow.python.autograph.pyct import parser
+from tensorflow.python.autograph.pyct import qual_names
+from tensorflow.python.autograph.pyct import templates
+from tensorflow.python.autograph.pyct.static_analysis import activity
+from tensorflow.python.autograph.pyct.static_analysis.annos import NodeAnno
+
+
+class _Statement(object):
+
+  def __init__(self):
+    self.pop_uses = None
+
+
+class ListTransformer(converter.Base):
+  """Converts lists and related operations to their TF counterpart."""
+
+  def visit_List(self, node):
+    node = self.generic_visit(node)
+    template = """
+      ag__.new_list(elements)
+    """
+    return templates.replace_as_expression(template, elements=node)
+
+  def _replace_append_call(self, node):
+    assert len(node.args) == 1
+    assert isinstance(node.func, gast.Attribute)
+    template = """
+      target = ag__.list_append(target, element)
+    """
+    return templates.replace(
+        template,
+        target=node.func.value,
+        element=node.args[0])
+
+  def _replace_pop_call(self, node):
+    # Expressions that use pop() are converted to a statement + expression.
+    #
+    # For example:
+    #
+    #   print(target.pop())
+    #
+    # ... is converted to:
+    #
+    #   target, target_pop = ag__.list_pop(target)
+    #   print(target_pop)
+    #
+    # Here, we just generate the variable name and swap it in,
+    # and _generate_pop_operation will handle the rest.
+    #
+    # Multiple uses of pop() are allowed:
+    #
+    #   print(tartget.pop(), target.pop())
+    #   print(tartget.pop().pop())
+    #
+    assert isinstance(node.func, gast.Attribute)
+    scope = anno.getanno(node, NodeAnno.ARGS_SCOPE)
+    target_node = node.func.value
+
+    # Attempt to use a related name if one exists. Otherwise use something
+    # generic.
+    if anno.hasanno(target_node, anno.Basic.QN):
+      target_name = anno.getanno(target_node, anno.Basic.QN).ssf()
+    else:
+      target_name = 'list_'
+    pop_var_name = self.ctx.namer.new_symbol(target_name, scope.referenced)
+
+    stmt = self.state[_Statement]
+    if stmt.pop_uses is None:
+      stmt.pop_uses = []
+    stmt.pop_uses.append((node, pop_var_name))
+
+    return templates.replace_as_expression('var_name', var_name=pop_var_name)
+
+  def _replace_stack_call(self, node):
+    assert len(node.args) == 1
+    dtype = self.get_definition_directive(
+        node.args[0],
+        directives.set_element_type,
+        'dtype',
+        default=templates.replace_as_expression('None'))
+    template = """
+      ag__.list_stack(
+          target,
+          opts=ag__.ListStackOpts(
+              element_dtype=dtype,
+              original_call=orig_call))
+    """
+    return templates.replace_as_expression(
+        template,
+        dtype=dtype,
+        target=node.args[0],
+        orig_call=node.func)
+
+  def visit_Call(self, node):
+    node = self.generic_visit(node)
+
+    # TODO(mdan): This is insufficient if target is a function argument.
+    # In the case of function arguments, we need to add the list to the
+    # function's return value, because it is being modified.
+    # TODO(mdan): Checking just the name is brittle, can it be improved?
+    if isinstance(node.func, gast.Attribute):
+      func_name = node.func.attr
+      if func_name == 'append' and (len(node.args) == 1):
+        node = self._replace_append_call(node)
+      elif func_name == 'pop' and (len(node.args) <= 1):
+        node = self._replace_pop_call(node)
+      elif (func_name == 'stack' and (len(node.args) == 1) and
+            (not node.keywords or node.keywords[0].arg == 'strict')):
+        # This avoids false positives with keyword args.
+        # TODO(mdan): handle kwargs properly.
+        node = self._replace_stack_call(node)
+
+    return node
+
+  def _generate_pop_operation(self, original_call_node, pop_var_name):
+    assert isinstance(original_call_node.func, gast.Attribute)
+
+    if original_call_node.args:
+      pop_element = original_call_node.args[0]
+    else:
+      pop_element = parser.parse_expression('None')
+
+    # The call will be something like "target.pop()", and the dtype is hooked to
+    # target, hence the func.value.
+    # TODO(mdan): For lists of lists, this won't work.
+    # The reason why it won't work is because it's unclear how to annotate
+    # the list as a "list of lists with a certain element type" when using
+    # operations like `l.pop().pop()`.
+    dtype = self.get_definition_directive(
+        original_call_node.func.value,
+        directives.set_element_type,
+        'dtype',
+        default=templates.replace_as_expression('None'))
+    shape = self.get_definition_directive(
+        original_call_node.func.value,
+        directives.set_element_type,
+        'shape',
+        default=templates.replace_as_expression('None'))
+
+    template = """
+      target, pop_var_name = ag__.list_pop(
+          target, element,
+          opts=ag__.ListPopOpts(element_dtype=dtype, element_shape=shape))
+    """
+    return templates.replace(
+        template,
+        target=original_call_node.func.value,
+        pop_var_name=pop_var_name,
+        element=pop_element,
+        dtype=dtype,
+        shape=shape)
+
+  def _postprocess_statement(self, node):
+    """Inserts any separate pop() calls that node may use."""
+    pop_uses = self.state[_Statement].pop_uses
+    if pop_uses:
+      replacements = []
+      for original_call_node, pop_var_name in pop_uses:
+        replacements.extend(
+            self._generate_pop_operation(original_call_node, pop_var_name))
+      replacements.append(node)
+      node = replacements
+    self.state[_Statement].exit()
+    return node, None
+
+  def _visit_and_process_block(self, block):
+    return self.visit_block(
+        block,
+        before_visit=self.state[_Statement].enter,
+        after_visit=self._postprocess_statement)
+
+  def visit_FunctionDef(self, node):
+    node.args = self.generic_visit(node.args)
+    node.decorator_list = self.visit_block(node.decorator_list)
+    node.body = self._visit_and_process_block(node.body)
+    return node
+
+  def visit_For(self, node):
+    node.target = self.visit(node.target)
+    node.body = self._visit_and_process_block(node.body)
+    node.orelse = self._visit_and_process_block(node.orelse)
+    return node
+
+  def visit_While(self, node):
+    node.test = self.visit(node.test)
+    node.body = self._visit_and_process_block(node.body)
+    node.orelse = self._visit_and_process_block(node.orelse)
+    return node
+
+  def visit_If(self, node):
+    node.test = self.visit(node.test)
+    node.body = self._visit_and_process_block(node.body)
+    node.orelse = self._visit_and_process_block(node.orelse)
+    return node
+
+  def visit_With(self, node):
+    node.items = self.visit_block(node.items)
+    node.body = self._visit_and_process_block(node.body)
+    return node
+
+
+def transform(node, ctx):
+  node = qual_names.resolve(node)
+  node = activity.resolve(node, ctx, None)
+
+  return ListTransformer(ctx).visit(node)

tensorflow/python/autograph/converters/lists_test.py

@@ -0,0 +1,107 @@
+# 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 lists module."""
+
+from tensorflow.python.autograph.converters import directives as directives_converter
+from tensorflow.python.autograph.converters import lists
+from tensorflow.python.autograph.core import converter_testing
+from tensorflow.python.autograph.lang import directives
+from tensorflow.python.autograph.lang import special_functions
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor
+from tensorflow.python.ops import array_ops_stack
+from tensorflow.python.ops import list_ops
+from tensorflow.python.platform import test
+
+
+class ListTest(converter_testing.TestCase):
+
+  def test_empty_list(self):
+
+    def f():
+      return []
+
+    tr = self.transform(f, lists)
+
+    tl = tr()
+    # Empty tensor lists cannot be evaluated or stacked.
+    self.assertIsInstance(tl, tensor.Tensor)
+    self.assertEqual(tl.dtype, dtypes.variant)
+
+  def test_initialized_list(self):
+
+    def f():
+      return [1, 2, 3]
+
+    tr = self.transform(f, lists)
+
+    self.assertAllEqual(tr(), [1, 2, 3])
+
+  def test_list_append(self):
+
+    def f():
+      l = special_functions.tensor_list([1])
+      l.append(2)
+      l.append(3)
+      return l
+
+    tr = self.transform(f, lists)
+
+    tl = tr()
+    r = list_ops.tensor_list_stack(tl, dtypes.int32)
+    self.assertAllEqual(self.evaluate(r), [1, 2, 3])
+
+  def test_list_pop(self):
+
+    def f():
+      l = special_functions.tensor_list([1, 2, 3])
+      directives.set_element_type(l, dtype=dtypes.int32, shape=())
+      s = l.pop()
+      return s, l
+
+    tr = self.transform(f, (directives_converter, lists))
+
+    ts, tl = tr()
+    r = list_ops.tensor_list_stack(tl, dtypes.int32)
+    self.assertAllEqual(self.evaluate(r), [1, 2])
+    self.assertAllEqual(self.evaluate(ts), 3)
+
+  def test_double_list_pop(self):
+
+    def f(l):
+      s = l.pop().pop()
+      return s
+
+    tr = self.transform(f, lists)
+
+    test_input = [1, 2, [1, 2, 3]]
+    # TODO(mdan): Pass a list of lists of tensor when we fully support that.
+    # For now, we just pass a regular Python list of lists just to verify that
+    # the two pop calls are sequenced properly.
+    self.assertAllEqual(tr(test_input), 3)
+
+  def test_list_stack(self):
+
+    def f():
+      l = [1, 2, 3]
+      return array_ops_stack.stack(l)
+
+    tr = self.transform(f, lists)
+
+    self.assertAllEqual(self.evaluate(tr()), [1, 2, 3])
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/autograph/converters/logical_expressions.py

@@ -0,0 +1,131 @@
+# Copyright 2016 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.
+# ==============================================================================
+"""Converter for logical expressions, e.g. `a and b -> tf.logical_and(a, b)`."""
+
+import gast
+
+from tensorflow.python.autograph.core import converter
+from tensorflow.python.autograph.pyct import parser
+from tensorflow.python.autograph.pyct import templates
+
+# TODO(mdan): Properly extract boolean ops according to lazy eval rules.
+# Note that this isn't completely safe either, because tensors may have control
+# dependencies.
+# Note that for loops that should be done after the loop was converted to
+# tf.while_loop so that the expanded conditionals are properly scoped.
+
+# Used to signal that an operand is safe for non-lazy evaluation.
+SAFE_BOOLEAN_OPERAND = 'SAFE_BOOLEAN_OPERAND'
+
+
+LOGICAL_OPERATORS = {
+    gast.And: 'ag__.and_',
+    gast.Not: 'ag__.not_',
+    gast.Or: 'ag__.or_',
+}
+
+EQUALITY_OPERATORS = {
+    gast.Eq: 'ag__.eq',
+    gast.NotEq: 'ag__.not_eq',
+}
+
+
+class LogicalExpressionTransformer(converter.Base):
+  """Converts logical expressions to corresponding TF calls."""
+
+  def _overload_of(self, operator):
+    op_type = type(operator)
+    if op_type in LOGICAL_OPERATORS:
+      return LOGICAL_OPERATORS[op_type]
+    if self.ctx.user.options.uses(converter.Feature.EQUALITY_OPERATORS):
+      if op_type in EQUALITY_OPERATORS:
+        return EQUALITY_OPERATORS[op_type]
+    return None
+
+  def _as_lambda(self, expr):
+    return templates.replace_as_expression('lambda: expr', expr=expr)
+
+  def _as_binary_function(self, func_name, arg1, arg2):
+    return templates.replace_as_expression(
+        'func_name(arg1, arg2)',
+        func_name=parser.parse_expression(func_name),
+        arg1=arg1,
+        arg2=arg2)
+
+  def _as_binary_operation(self, op, arg1, arg2):
+    template = templates.replace_as_expression(
+        'arg1 is arg2',  # Note: `is` will be replaced with `op` below.
+        arg1=arg1,
+        arg2=arg2)
+    template.ops[0] = op
+    return template
+
+  def _as_unary_function(self, func_name, arg):
+    return templates.replace_as_expression(
+        'func_name(arg)', func_name=parser.parse_expression(func_name), arg=arg)
+
+  def _process_binop(self, op, left, right):
+    overload = self._overload_of(op)
+    if overload is None:
+      return self._as_binary_operation(op, left, right)
+    return self._as_binary_function(overload, left, right)
+
+  def visit_Compare(self, node):
+    node = self.generic_visit(node)
+
+    ops_and_comps = list(zip(node.ops, node.comparators))
+    left = node.left
+
+    # Repeated comparisons are converted to conjunctions:
+    #   a < b < c   ->   a < b and b < c
+    op_tree = None
+    while ops_and_comps:
+      op, right = ops_and_comps.pop(0)
+      binary_comparison = self._process_binop(op, left, right)
+      if op_tree is not None:
+        op_tree = self._as_binary_function('ag__.and_',
+                                           self._as_lambda(op_tree),
+                                           self._as_lambda(binary_comparison))
+      else:
+        op_tree = binary_comparison
+      left = right
+
+    assert op_tree is not None
+    return op_tree
+
+  def visit_UnaryOp(self, node):
+    node = self.generic_visit(node)
+
+    overload = self._overload_of(node.op)
+    if overload is None:
+      return node
+
+    return self._as_unary_function(overload, node.operand)
+
+  def visit_BoolOp(self, node):
+    node = self.generic_visit(node)
+    node_values = node.values
+    right = node.values.pop()
+    while node_values:
+      left = node_values.pop()
+      right = self._as_binary_function(
+          self._overload_of(node.op), self._as_lambda(left),
+          self._as_lambda(right))
+    return right
+
+
+def transform(node, ctx):
+  transformer = LogicalExpressionTransformer(ctx)
+  return transformer.visit(node)

tensorflow/python/autograph/converters/logical_expressions_test.py

@@ -0,0 +1,82 @@
+# 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 logical_expressions module."""
+
+from tensorflow.python.autograph.converters import logical_expressions
+from tensorflow.python.autograph.core import converter_testing
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import test_util
+from tensorflow.python.platform import test
+
+
+class LogicalExpressionTest(converter_testing.TestCase):
+
+  def test_equals(self):
+
+    def f(a, b):
+      return a == b
+
+    tr = self.transform(f, logical_expressions)
+
+    self.assertTrue(self.evaluate(tr(constant_op.constant(1), 1)))
+    self.assertFalse(self.evaluate(tr(constant_op.constant(1), 2)))
+
+  @test_util.run_deprecated_v1
+  def test_bool_ops(self):
+
+    def f(a, b, c):
+      return (a or b) and (a or b or c) and not c
+
+    tr = self.transform(f, logical_expressions)
+
+    self.assertTrue(self.evaluate(tr(constant_op.constant(True), False, False)))
+    self.assertFalse(self.evaluate(tr(constant_op.constant(True), False, True)))
+
+  def test_comparison(self):
+
+    def f(a, b, c, d):
+      return a < b == c > d
+
+    tr = self.transform(f, logical_expressions)
+
+    # Note: having just the first constant a tensor tests that the
+    # operations execute in the correct order. If anything other than
+    # a < b executed first, the result would be a Python scalar and not a
+    # Tensor. This is valid as long as the dispat is automatic based on
+    # type.
+    self.assertTrue(self.evaluate(tr(constant_op.constant(1), 2, 2, 1)))
+    self.assertFalse(self.evaluate(tr(constant_op.constant(1), 2, 2, 3)))
+
+  def test_default_ops(self):
+
+    def f(a, b):
+      return a in b
+
+    tr = self.transform(f, logical_expressions)
+
+    self.assertTrue(tr('a', ('a',)))
+
+  def test_unary_ops(self):
+
+    def f(a):
+      return ~a, -a, +a
+
+    tr = self.transform(f, logical_expressions)
+
+    self.assertEqual(tr(1), (-2, -1, 1))
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/autograph/converters/return_statements.py

@@ -0,0 +1,402 @@
+# 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.
+# ==============================================================================
+"""Canonicalizes functions with multiple returns to use just one."""
+
+import gast
+
+from tensorflow.python.autograph.core import converter
+from tensorflow.python.autograph.pyct import anno
+from tensorflow.python.autograph.pyct import parser
+from tensorflow.python.autograph.pyct import qual_names
+from tensorflow.python.autograph.pyct import templates
+from tensorflow.python.autograph.pyct.static_analysis import activity
+from tensorflow.python.autograph.pyct.static_analysis.annos import NodeAnno
+
+
+BODY_DEFINITELY_RETURNS = 'BODY_DEFINITELY_RETURNS'
+ORELSE_DEFINITELY_RETURNS = 'ORELSE_DEFINITELY_RETURNS'
+STMT_DEFINITELY_RETURNS = 'STMT_DEFINITELY_RETURNS'
+
+
+class _RewriteBlock(object):
+
+  def __init__(self):
+    self.definitely_returns = False
+
+
+class ConditionalReturnRewriter(converter.Base):
+  """Rewrites a pattern where it's unobvious that all paths return a value.
+
+  This rewrite allows avoiding intermediate None return values.
+
+  The following pattern:
+
+      if cond:
+        <block 1>
+        return
+      else:
+        <block 2>
+      <block 3>
+
+  is converted to:
+
+      if cond:
+        <block 1>
+        return
+      else:
+        <block 2>
+        <block 3>
+
+  and vice-versa (if the else returns, subsequent statements are moved under the
+  if branch).
+  """
+
+  def visit_Return(self, node):
+    self.state[_RewriteBlock].definitely_returns = True
+    return node
+
+  def _postprocess_statement(self, node):
+    # If the node definitely returns (e.g. it's a with statement with a
+    # return statement in it), then the current block also definitely returns.
+    if anno.getanno(node, STMT_DEFINITELY_RETURNS, default=False):
+      self.state[_RewriteBlock].definitely_returns = True
+
+    # The special case: collapse a typical conditional return pattern into
+    # a single conditional with possibly returns on both branches. This
+    # reduces the use of None return values, which don't work with TF
+    # conditionals.
+    if (isinstance(node, gast.If)
+        and anno.getanno(node, BODY_DEFINITELY_RETURNS, default=False)):
+      return node, node.orelse
+    elif (isinstance(node, gast.If)
+          and anno.getanno(node, ORELSE_DEFINITELY_RETURNS, default=False)):
+      return node, node.body
+
+    return node, None
+
+  def _visit_statement_block(self, node, nodes):
+    self.state[_RewriteBlock].enter()
+    new_nodes = self.visit_block(nodes, after_visit=self._postprocess_statement)
+    block_definitely_returns = self.state[_RewriteBlock].definitely_returns
+    self.state[_RewriteBlock].exit()
+    return new_nodes, block_definitely_returns
+
+  def visit_While(self, node):
+    node.test = self.visit(node.test)
+    node.body, _ = self._visit_statement_block(node, node.body)
+    node.orelse, _ = self._visit_statement_block(node, node.orelse)
+    return node
+
+  def visit_For(self, node):
+    node.iter = self.visit(node.iter)
+    node.target = self.visit(node.target)
+    node.body, _ = self._visit_statement_block(node, node.body)
+    node.orelse, _ = self._visit_statement_block(node, node.orelse)
+    return node
+
+  def visit_With(self, node):
+    node.items = self.visit_block(node.items)
+    node.body, definitely_returns = self._visit_statement_block(node, node.body)
+    if definitely_returns:
+      anno.setanno(node, STMT_DEFINITELY_RETURNS, True)
+    return node
+
+  def visit_Try(self, node):
+    # We could decide whether a 'try' DEFINITELY_RETURNS based on its components
+    # It is not clear whether we want to do anything with this given
+    # a 'try' is likely to throw an exception in some circumstances.
+    node.body, _ = self._visit_statement_block(node, node.body)
+    node.orelse, _ = self._visit_statement_block(node, node.orelse)
+    node.finalbody, _ = self._visit_statement_block(node, node.finalbody)
+    node.handlers = self.visit_block(node.handlers)
+    return node
+
+  def visit_ExceptHandler(self, node):
+    # To determine whether `try` DEFINITELY_RETURNS we need to revisit this.
+    node.body, _ = self._visit_statement_block(node, node.body)
+    return node
+
+  def visit_If(self, node):
+    node.test = self.visit(node.test)
+
+    node.body, body_definitely_returns = self._visit_statement_block(
+        node, node.body)
+    if body_definitely_returns:
+      anno.setanno(node, BODY_DEFINITELY_RETURNS, True)
+
+    node.orelse, orelse_definitely_returns = self._visit_statement_block(
+        node, node.orelse)
+    if orelse_definitely_returns:
+      anno.setanno(node, ORELSE_DEFINITELY_RETURNS, True)
+
+    if body_definitely_returns and orelse_definitely_returns:
+      self.state[_RewriteBlock].definitely_returns = True
+
+    return node
+
+  def visit_FunctionDef(self, node):
+    node.args = self.visit(node.args)
+    node.body, _ = self._visit_statement_block(node, node.body)
+    return node
+
+
+class _Block(object):
+
+  def __init__(self):
+    self.is_function = False
+    self.return_used = False
+    self.create_guard_next = False
+    self.create_guard_now = False
+
+  def __repr__(self):
+    return 'used: {}'.format(
+        self.return_used)
+
+
+class _Function(object):
+
+  def __init__(self):
+    self.do_return_var_name = None
+    self.retval_var_name = None
+
+  def __repr__(self):
+    return 'return control: {}, return value: {}'.format(
+        self.do_return_var_name, self.retval_var_name)
+
+
+class ReturnStatementsTransformer(converter.Base):
+  """Lowers return statements into variables and conditionals.
+
+  Specifically, the following pattern:
+
+      <block 1>
+      return val
+      <block 2>
+
+  is converted to:
+
+      do_return = False
+      retval = None
+
+      <block 1>
+
+      do_return = True
+      retval = val
+
+      if not do_return:
+        <block 2>
+
+      return retval
+
+  The conversion adjusts loops as well:
+
+      <block 1>
+      while cond:
+        <block 2>
+        return retval
+
+  is converted to:
+
+      <block 1>
+      while not do_return and cond:
+        <block 2>
+        do_return = True
+        retval = val
+  """
+
+  def __init__(self, ctx, allow_missing_return):
+    super(ReturnStatementsTransformer, self).__init__(ctx)
+    self.allow_missing_return = allow_missing_return
+
+  def visit_Return(self, node):
+    for block in reversed(self.state[_Block].stack):
+      block.return_used = True
+      block.create_guard_next = True
+      if block.is_function:
+        break
+
+    retval = node.value if node.value else parser.parse_expression('None')
+
+    # Note: If `return <expr> raises, then the return is aborted.
+    # The try-catch below ensures the variables remain consistent in that case.
+    template = """
+      try:
+        do_return_var_name = True
+        retval_var_name = retval
+      except:
+        do_return_var_name = False
+        raise
+    """
+    node = templates.replace(
+        template,
+        do_return_var_name=self.state[_Function].do_return_var_name,
+        retval_var_name=self.state[_Function].retval_var_name,
+        retval=retval)
+
+    return node
+
+  def _postprocess_statement(self, node):
+    if not self.state[_Block].return_used:
+      return node, None
+
+    state = self.state[_Block]
+    if state.create_guard_now:
+      template = """
+        if not do_return_var_name:
+          original_node
+      """
+      cond, = templates.replace(
+          template,
+          do_return_var_name=self.state[_Function].do_return_var_name,
+          original_node=node)
+      node, block = cond, cond.body
+    else:
+      node, block = node, None
+
+    state.create_guard_now = state.create_guard_next
+    state.create_guard_next = False
+
+    return node, block
+
+  def _visit_statement_block(self, node, nodes):
+    self.state[_Block].enter()
+    nodes = self.visit_block(nodes, after_visit=self._postprocess_statement)
+    self.state[_Block].exit()
+    return nodes
+
+  def visit_While(self, node):
+    node.test = self.visit(node.test)
+
+    # Add the check for return to the loop condition.
+    node.body = self._visit_statement_block(node, node.body)
+    if self.state[_Block].return_used:
+      node.test = templates.replace_as_expression(
+          'not control_var and test',
+          test=node.test,
+          control_var=self.state[_Function].do_return_var_name)
+
+    node.orelse = self._visit_statement_block(node, node.orelse)
+    return node
+
+  def visit_For(self, node):
+    node.iter = self.visit(node.iter)
+    node.target = self.visit(node.target)
+
+    # Add the check for return to the loop condition.
+    node.body = self._visit_statement_block(node, node.body)
+    if self.state[_Block].return_used:
+      extra_test = anno.getanno(node, anno.Basic.EXTRA_LOOP_TEST, default=None)
+      if extra_test is not None:
+        extra_test = templates.replace_as_expression(
+            'not control_var and extra_test',
+            extra_test=extra_test,
+            control_var=self.state[_Function].do_return_var_name)
+      else:
+        extra_test = templates.replace_as_expression(
+            'not control_var',
+            control_var=self.state[_Function].do_return_var_name)
+      anno.setanno(node, anno.Basic.EXTRA_LOOP_TEST, extra_test)
+
+    node.orelse = self._visit_statement_block(node, node.orelse)
+    return node
+
+  def visit_With(self, node):
+    node.items = self.visit_block(node.items)
+    node.body = self._visit_statement_block(node, node.body)
+    return node
+
+  def visit_Try(self, node):
+    node.body = self._visit_statement_block(node, node.body)
+    node.orelse = self._visit_statement_block(node, node.orelse)
+    node.finalbody = self._visit_statement_block(node, node.finalbody)
+    node.handlers = self.visit_block(node.handlers)
+    return node
+
+  def visit_ExceptHandler(self, node):
+    node.body = self._visit_statement_block(node, node.body)
+    return node
+
+  def visit_If(self, node):
+    node.test = self.visit(node.test)
+    node.body = self._visit_statement_block(node, node.body)
+    node.orelse = self._visit_statement_block(node, node.orelse)
+    return node
+
+  def visit_FunctionDef(self, node):
+    with self.state[_Function] as fn:
+      with self.state[_Block] as block:
+        block.is_function = True
+
+        scope = anno.getanno(node, NodeAnno.BODY_SCOPE)
+        do_return_var_name = self.ctx.namer.new_symbol('do_return',
+                                                       scope.referenced)
+        retval_var_name = self.ctx.namer.new_symbol('retval_', scope.referenced)
+        fn.do_return_var_name = do_return_var_name
+        fn.retval_var_name = retval_var_name
+
+        node.body = self._visit_statement_block(node, node.body)
+
+        if block.return_used:
+
+          if self.allow_missing_return:
+            # The function would have a single `with` node that wraps the
+            # entire body. If the function had a docstring, the body has two
+            # nodes, with the `with` as the second node.
+            wrapper_node = node.body[-1]
+            assert isinstance(wrapper_node, gast.With), (
+                'This transformer requires the functions converter.')
+
+            template = """
+              do_return_var_name = False
+              retval_var_name = ag__.UndefinedReturnValue()
+              body
+              return function_context.ret(retval_var_name, do_return_var_name)
+            """
+
+            wrapper_node.body = templates.replace(
+                template,
+                body=wrapper_node.body,
+                do_return_var_name=do_return_var_name,
+                function_context=anno.getanno(node, 'function_context_name'),
+                retval_var_name=retval_var_name)
+          else:
+            template = """
+              body
+              return retval_var_name
+            """
+            node.body = templates.replace(
+                template,
+                body=node.body,
+                do_return_var_name=do_return_var_name,
+                retval_var_name=retval_var_name)
+
+    return node
+
+
+def transform(node, ctx, default_to_null_return=True):
+  """Ensure a function has only a single return, at the end."""
+  node = qual_names.resolve(node)
+  node = activity.resolve(node, ctx, None)
+
+  # Note: Technically, these two could be merged into a single walk, but
+  # keeping them separate helps with readability.
+  node = ConditionalReturnRewriter(ctx).visit(node)
+
+  node = qual_names.resolve(node)
+  node = activity.resolve(node, ctx, None)
+  transformer = ReturnStatementsTransformer(
+      ctx, allow_missing_return=default_to_null_return)
+  node = transformer.visit(node)
+  return node

tensorflow/python/autograph/converters/return_statements_test.py

@@ -0,0 +1,254 @@
+# 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 return_statements module."""
+
+from tensorflow.python.autograph.converters import functions
+from tensorflow.python.autograph.converters import return_statements
+from tensorflow.python.autograph.core import converter_testing
+from tensorflow.python.framework import ops
+from tensorflow.python.platform import test
+
+
+class SingleReturnTest(converter_testing.TestCase):
+
+  def assertTransformedEquivalent(self, f, *inputs):
+    tr = self.transform(f, (functions, return_statements))
+    self.assertEqual(f(*inputs), tr(*inputs))
+
+  def test_straightline(self):
+
+    def f(x):
+      return x * x
+
+    self.assertTransformedEquivalent(f, 2)
+
+  def test_superfluous_returns(self):
+
+    def f():
+      retval = 1
+      return retval
+      retval = 2  # pylint:disable=unreachable
+      return retval
+
+    self.assertTransformedEquivalent(f)
+
+  def test_superfluous_returns_adjacent(self):
+
+    def f():
+      return 1
+      return 2  # pylint:disable=unreachable
+
+    self.assertTransformedEquivalent(f)
+
+  def test_conditional(self):
+
+    def f(x):
+      if x > 0:
+        return x
+      else:
+        return x * x
+
+    self.assertTransformedEquivalent(f, 2)
+    self.assertTransformedEquivalent(f, -2)
+
+  def test_conditional_missing_else(self):
+
+    def f(x):
+      if x > 0:
+        return x
+
+    self.assertTransformedEquivalent(f, 2)
+    self.assertTransformedEquivalent(f, -2)
+
+  def test_conditional_missing_else_then_default(self):
+
+    def f(x):
+      if x > 0:
+        return x
+      return x * x
+
+    self.assertTransformedEquivalent(f, 2)
+    self.assertTransformedEquivalent(f, -2)
+
+  def test_conditional_else_only_then_default(self):
+
+    def f(x):
+      if x < 0:
+        x *= x
+      else:
+        return x
+      return x
+
+    self.assertTransformedEquivalent(f, 2)
+    self.assertTransformedEquivalent(f, -2)
+
+  def test_conditional_nested(self):
+
+    def f(x):
+      if x > 0:
+        if x < 5:
+          return x
+        else:
+          return x * x
+      else:
+        return x * x * x
+
+    self.assertTransformedEquivalent(f, 2)
+    self.assertTransformedEquivalent(f, -2)
+    self.assertTransformedEquivalent(f, 5)
+
+  def test_context_manager(self):
+
+    def f(x):
+      with ops.name_scope(''):
+        return x * x
+
+    self.assertTransformedEquivalent(f, 2)
+    self.assertTransformedEquivalent(f, -2)
+
+  def test_context_manager_in_conditional(self):
+
+    def f(x):
+      if x > 0:
+        with ops.name_scope(''):
+          return x * x
+      else:
+        return x
+
+    self.assertTransformedEquivalent(f, 2)
+    self.assertTransformedEquivalent(f, -2)
+
+  def test_conditional_in_context_manager(self):
+
+    def f(x):
+      with ops.name_scope(''):
+        if x > 0:
+          return x * x
+        else:
+          return x
+
+    self.assertTransformedEquivalent(f, 2)
+    self.assertTransformedEquivalent(f, -2)
+
+  def test_no_return(self):
+
+    def f(x):
+      x *= x
+
+    self.assertTransformedEquivalent(f, 2)
+
+  def test_nested_function(self):
+
+    def f(x):
+
+      def inner_fn(y):
+        if y > 0:
+          return y * y
+        else:
+          return y
+
+      return inner_fn(x)
+
+    self.assertTransformedEquivalent(f, 2)
+    self.assertTransformedEquivalent(f, -2)
+
+  def test_nested_function_in_control_flow(self):
+
+    def f(x):
+
+      if x:
+        def inner_fn(y):
+          return y
+        inner_fn(x)
+
+    self.assertTransformedEquivalent(f, 2)
+    self.assertTransformedEquivalent(f, -2)
+
+  def test_for_loop(self):
+
+    def f(n):
+      for _ in range(n):
+        return 1
+
+    self.assertTransformedEquivalent(f, 2)
+    self.assertTransformedEquivalent(f, 0)
+
+  def test_while_loop(self):
+
+    def f(n):
+      i = 0
+      s = 0
+      while i < n:
+        i += 1
+        s += i
+        if s > 4:
+          return s
+      return -1
+
+    self.assertTransformedEquivalent(f, 0)
+    self.assertTransformedEquivalent(f, 2)
+    self.assertTransformedEquivalent(f, 4)
+
+  def test_null_return(self):
+
+    def f(n):
+      if n > 4:
+        return
+      return
+
+    self.assertTransformedEquivalent(f, 4)
+    self.assertTransformedEquivalent(f, 5)
+
+  def test_nested_multiple_withs(self):
+
+    def f(x):
+      v = []
+      while x > 0:
+        x -= 1
+        with ops.name_scope(''):
+          if x % 2 == 0:
+            return v
+        with ops.name_scope(''):
+          v.append(x)
+        v.append(x)
+      return v
+
+    self.assertTransformedEquivalent(f, 0)
+    self.assertTransformedEquivalent(f, 1)
+    self.assertTransformedEquivalent(f, 3)
+    self.assertTransformedEquivalent(f, 4)
+
+  def test_multiple_returns_in_nested_scope(self):
+
+    def f(a):
+      v = []
+      for x in a:
+        x -= 1
+        if x > 100:
+          return v
+        try:
+          raise ValueError('intentional')
+        except ValueError:  # pylint:disable=bare-except
+          return v
+        v.append(x)
+      return v
+
+    self.assertTransformedEquivalent(f, [])
+    self.assertTransformedEquivalent(f, [1])
+    self.assertTransformedEquivalent(f, [2])
+    self.assertTransformedEquivalent(f, [1, 2, 3])
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/autograph/converters/slices.py

@@ -0,0 +1,83 @@
+# Copyright 2016 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.
+# ==============================================================================
+"""Converter for slice operations."""
+
+import gast
+
+from tensorflow.python.autograph.core import converter
+from tensorflow.python.autograph.lang import directives
+from tensorflow.python.autograph.pyct import templates
+
+
+class SliceTransformer(converter.Base):
+  """Converts slicing operations to their TF counterpart.
+
+  Currently, relying on the default slice operator that Tensor uses is
+  insufficient, because TensorArray and tensor lists use dedicated index read
+  and write functions.
+  """
+
+  def _process_single_assignment(self, target, value):
+    if not isinstance(target, gast.Subscript):
+      return None
+    s = target.slice
+    if isinstance(s, (gast.Tuple, gast.Slice)):
+      return None
+
+    template = """
+      target = ag__.set_item(target, key, item)
+    """
+    return templates.replace(
+        template, target=target.value, key=target.slice, item=value)
+
+  def visit_Assign(self, node):
+    node = self.generic_visit(node)
+    # TODO(mdan): Support unpackings and multiple assignments.
+    if len(node.targets) != 1:
+      raise NotImplementedError('multiple assignment')
+    replacement = self._process_single_assignment(node.targets[0], node.value)
+    if replacement is not None:
+      return replacement
+    return node
+
+  def visit_Subscript(self, node):
+    node = self.generic_visit(node)
+    s = node.slice
+    if isinstance(s, (gast.Tuple, gast.Slice)):
+      return node
+
+    if not isinstance(node.ctx, gast.Load):
+      # Index writes are handled at a higher level, one at which the rvalue is
+      # also available.
+      return node
+
+    dtype = self.get_definition_directive(
+        node.value,
+        directives.set_element_type,
+        'dtype',
+        default=templates.replace_as_expression('None'))
+
+    template = """
+      ag__.get_item(
+          target,
+          key,
+          opts=ag__.GetItemOpts(element_dtype=dtype))
+    """
+    return templates.replace_as_expression(
+        template, target=node.value, key=s, dtype=dtype)
+
+
+def transform(node, ctx):
+  return SliceTransformer(ctx).visit(node)

tensorflow/python/autograph/converters/slices_test.py

@@ -0,0 +1,54 @@
+# 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 slices module."""
+
+from tensorflow.python.autograph.converters import directives as directives_converter
+from tensorflow.python.autograph.converters import slices
+from tensorflow.python.autograph.core import converter_testing
+from tensorflow.python.autograph.lang import directives
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.ops import list_ops
+from tensorflow.python.platform import test
+
+
+class SliceTest(converter_testing.TestCase):
+
+  def test_index_access(self):
+
+    def f(l):
+      directives.set_element_type(l, dtypes.int32)
+      return l[1]
+
+    tr = self.transform(f, (directives_converter, slices))
+
+    tl = list_ops.tensor_list_from_tensor(
+        [1, 2], element_shape=constant_op.constant([], dtype=dtypes.int32))
+    y = tr(tl)
+    self.assertEqual(2, self.evaluate(y))
+
+  def test_index_access_multiple_definitions(self):
+
+    def f(l):
+      directives.set_element_type(l, dtypes.int32)
+      if l:
+        l = []
+      return l[1]
+
+    self.transform(f, (directives_converter, slices))
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/autograph/converters/variables.py

@@ -0,0 +1,97 @@
+# Copyright 2020 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.
+# ==============================================================================
+"""Overloads all variable read operations."""
+
+import gast
+
+from tensorflow.python.autograph.core import converter
+from tensorflow.python.autograph.pyct import anno
+from tensorflow.python.autograph.pyct import templates
+
+
+class VariableAccessTransformer(converter.Base):
+  """Rewrites basic symbol reads.
+
+  This transformer rewrites variable reads with a "read" operator which allows
+  tracking activity.
+
+  Example:
+
+  For a basic statement:
+
+      a = b + c
+
+  This is translated to:
+
+      a = ld(b) + ld(c)
+
+  Augmented assignment operations also introduce a `ld` operator:
+
+      a += b
+
+  The assignment target also receives an operator to properly represent the
+  read:
+
+      a = ld(a)
+      a += ld(b)
+  """
+
+  def visit_Name(self, node):
+    # Only the loads which existed in the original code are overloaded.
+    if not anno.hasanno(node, anno.Static.ORIG_DEFINITIONS):
+      return node
+    if isinstance(node.ctx, gast.Load):
+      node = templates.replace_as_expression('ag__.ld(var_)', var_=node)
+    return node
+
+  def visit_Delete(self, node):
+    node = self.generic_visit(node)
+
+    rewrite_targets = []
+    for tgt in node.targets:
+      # Don't rewrite composites like `del a[0]`.
+      if isinstance(tgt, gast.Name):
+        rewrite_targets.append(tgt)
+
+    if not rewrite_targets:
+      return node
+
+    results = []
+    for tgt in rewrite_targets:
+      template = """
+        var_ = ag__.Undefined(var_name)
+      """
+      results.extend(templates.replace(
+          template, var_=tgt, var_name=gast.Constant(tgt.id, kind=None)))
+    remaining_targets = [n for n in node.targets if n not in rewrite_targets]
+    if remaining_targets:
+      results.append(gast.Delete(targets=remaining_targets))
+
+    return results
+
+  def visit_AugAssign(self, node):
+    if isinstance(node.target, gast.Name):
+      template = """
+        var_ = ag__.ld(var_)
+        original
+      """
+      node = templates.replace(template, var_=node.target, original=node)
+    else:
+      node = self.generic_visit(node)
+    return node
+
+
+def transform(node, ctx):
+  return VariableAccessTransformer(ctx).visit(node)

tensorflow/python/autograph/converters/variables_test.py

@@ -0,0 +1,210 @@
+# Copyright 2020 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 variables module."""
+
+from tensorflow.python.autograph.converters import variables
+from tensorflow.python.autograph.core import converter_testing
+from tensorflow.python.platform import test
+
+
+class VariablesTest(converter_testing.TestCase):
+
+  def transform_with_test_ld(self, f):
+    """Generates code which adds 1 to all variable reads."""
+    return self.transform(f, variables, ag_overrides={'ld': lambda x: x + 1})
+
+  def test_read(self):
+
+    def f(l):
+      return l
+
+    tr = self.transform_with_test_ld(f)
+
+    self.assertEqual(tr(1), 2)
+
+  def test_aug_assign(self):
+
+    def f(l):
+      l *= 10
+      return l
+
+    tr = self.transform_with_test_ld(f)
+
+    self.assertEqual(tr(1), (1 + 1) * 10 + 1)  # two reads
+
+  def test_del(self):
+
+    def f(l):
+      del l
+      return l
+
+    tr = self.transform(f, variables)
+
+    with self.assertRaisesRegex(NameError, "'l' is used before assignment"):
+      tr(1)
+
+  def test_del_getitem_ignored_basic_slice(self):
+
+    def f(l):
+      del l[0]
+      return l
+
+    tr = self.transform(f, variables)
+
+    self.assertListEqual([2], tr([1, 2]))
+
+  def test_del_getitem_ignored_range_slice(self):
+
+    def f(l):
+      del l[0:2]
+      return l
+
+    tr = self.transform(f, variables)
+
+    self.assertListEqual([], tr([1, 2]))
+
+  def test_del_getattr_ignored(self):
+
+    def f(l):
+      del l.a
+      return l
+
+    class TestClass(object):
+
+      def __init__(self):
+        self.a = 1
+        self.b = 2
+
+    tr = self.transform(f, variables)
+
+    self.assertFalse(hasattr(tr(TestClass()), 'a'))
+    self.assertEqual(tr(TestClass()).b, 2)
+
+  def test_del_packing_ignored_list(self):
+    # Note: testing for UnboundLocalError, not NameError because in this case we
+    # don't rewrite the del.
+
+    def f(a, b):
+      del [a, b]
+      return a
+
+    tr = self.transform(f, variables)
+
+    with self.assertRaises(UnboundLocalError):
+      tr(1, 2)
+
+  def test_del_packing_ignored_nested(self):
+    # Note: testing for UnboundLocalError, not NameError because in this case we
+    # don't rewrite the del.
+
+    def f(a, b, c):
+      del [a, (b, c)]
+      return c
+
+    tr = self.transform(f, variables)
+
+    with self.assertRaises(UnboundLocalError):
+      tr(1, 2, 3)
+
+  def test_del_item_multiple_mixed_used_after(self):
+
+    def f(a, b, c):
+      del a, b, c[0]
+      a = 1
+      return a, b, c
+
+    tr = self.transform(f, variables)
+
+    with self.assertRaisesRegex(NameError, "'b' is used before assignment"):
+      tr(1, 2, [1, 2])
+
+  def test_del_item_multiple_mixed_unused_after(self):
+
+    def f(a, b, c):
+      del a, b, c[0]
+      a = 1
+      b = 2
+      return c
+
+    tr = self.transform(f, variables)
+
+    self.assertListEqual([2], tr(1, 2, [1, 2]))
+
+  def test_attribute(self):
+
+    class TestClass(object):
+
+      def __init__(self):
+        self.v = 1
+
+      def __add__(self, other):
+        self.v += other
+        return self
+
+    def f(l):
+      return l.v
+
+    tc = TestClass()
+    tr = self.transform_with_test_ld(f)
+
+    self.assertEqual(tr(tc), 2)
+
+  def test_subscript(self):
+
+    class TestClass(object):
+
+      def __init__(self):
+        self.v = 1
+
+      def __add__(self, other):
+        self.v += other
+        return self
+
+      def __getitem__(self, _):
+        return self.v
+
+    def f(l):
+      return l[0]
+
+    tc = TestClass()
+    tr = self.transform_with_test_ld(f)
+
+    self.assertEqual(tr(tc), 2)
+
+  def test_call(self):
+
+    class TestClass(object):
+
+      def __init__(self):
+        self.v = 1
+
+      def __add__(self, other):
+        self.v += other
+        return self
+
+      def __call__(self):
+        return self.v
+
+    def f(l):
+      return l()
+
+    tc = TestClass()
+    tr = self.transform_with_test_ld(f)
+
+    self.assertEqual(tr(tc), 2)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/autograph/core/BUILD

@@ -0,0 +1,122 @@
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:tensorflow.bzl", "py_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "test_lib",
+    srcs = [
+        "converter_testing.py",
+    ],
+    strict_deps = True,
+    visibility = ["//tensorflow:__subpackages__"],
+    deps = [
+        ":config",
+        ":converter",
+        "//tensorflow/python/autograph/impl:api",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_library(
+    name = "ag_ctx",
+    srcs = ["ag_ctx.py"],
+    strict_deps = True,
+    visibility = ["//tensorflow:__subpackages__"],
+    deps = [
+        "//tensorflow/python/autograph/utils:ag_logging",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "unsupported_features_checker",
+    srcs = ["unsupported_features_checker.py"],
+    strict_deps = True,
+    visibility = ["//tensorflow:__subpackages__"],
+    deps = [
+        "//tensorflow/python/autograph/pyct:errors",
+        "@pypi//gast",
+    ],
+)
+
+py_library(
+    name = "config",
+    srcs = ["config.py"],
+    strict_deps = True,
+    visibility = ["//tensorflow:__subpackages__"],
+    deps = [":config_lib"],
+)
+
+py_library(
+    name = "config_lib",
+    srcs = ["config_lib.py"],
+    strict_deps = True,
+    visibility = ["//visibility:private"],
+)
+
+py_library(
+    name = "function_wrappers",
+    srcs = ["function_wrappers.py"],
+    strict_deps = True,
+    visibility = ["//tensorflow:__subpackages__"],
+    deps = [
+        ":ag_ctx",
+        ":converter",
+        "//tensorflow/python/autograph/operators:variables",
+        "//tensorflow/python/framework:auto_control_deps",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor_util",
+        "//tensorflow/python/util:nest",
+    ],
+)
+
+py_library(
+    name = "converter",
+    srcs = ["converter.py"],
+    strict_deps = True,
+    visibility = ["//tensorflow:__subpackages__"],
+    deps = [
+        "//tensorflow/python/autograph/pyct:anno",
+        "//tensorflow/python/autograph/pyct:ast_util",
+        "//tensorflow/python/autograph/pyct:parser",
+        "//tensorflow/python/autograph/pyct:templates",
+        "//tensorflow/python/autograph/pyct:transformer",
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_test(
+    name = "converter_test",
+    srcs = ["converter_test.py"],
+    strict_deps = True,
+    deps = [
+        ":converter",
+        ":test_lib",
+        #internal proto upb dep
+        "//tensorflow/python/autograph/pyct:anno",
+        "//tensorflow/python/autograph/pyct:loader",
+        "//tensorflow/python/autograph/pyct:parser",
+        "//tensorflow/python/autograph/pyct:templates",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_test(
+    name = "function_wrappers_test",
+    srcs = ["function_wrappers_test.py"],
+    strict_deps = True,
+    deps = [
+        ":converter",
+        ":function_wrappers",
+        #internal proto upb dep
+        "//tensorflow/python/eager:context",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)

tensorflow/python/autograph/core/ag_ctx.py

@@ -0,0 +1,105 @@
+# Copyright 2019 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.
+# ==============================================================================
+"""Thread-local context managers for AutoGraph."""
+
+import enum
+import inspect
+import threading
+
+from tensorflow.python.autograph.utils import ag_logging
+from tensorflow.python.util.tf_export import tf_export
+
+
+stacks = threading.local()
+
+
+def _control_ctx():
+  if not hasattr(stacks, 'control_status'):
+    stacks.control_status = [_default_control_status_ctx()]
+  return stacks.control_status
+
+
+@tf_export('__internal__.autograph.control_status_ctx', v1=[])
+def control_status_ctx():
+  """Returns the current control context for autograph.
+
+  This method is useful when calling `tf.__internal__.autograph.tf_convert`,
+  The context will be used by tf_convert to determine whether it should convert
+  the input function. See the sample usage like below:
+
+  ```
+  def foo(func):
+    return tf.__internal__.autograph.tf_convert(
+       input_fn, ctx=tf.__internal__.autograph.control_status_ctx())()
+  ```
+
+  Returns:
+    The current control context of autograph.
+  """
+  ret = _control_ctx()[-1]
+  return ret
+
+
+class Status(enum.Enum):
+  UNSPECIFIED = 0
+  ENABLED = 1
+  DISABLED = 2
+
+
+class ControlStatusCtx(object):
+  """A context that tracks whether autograph is enabled by the user."""
+
+  def __init__(self, status, options=None):
+    self.status = status
+    self.options = options
+
+  def __enter__(self):
+    _control_ctx().append(self)
+    return self
+
+  def __repr__(self):
+    return '{}[status={}, options={}]'.format(
+        self.__class__.__name__, self.status, self.options)
+
+  def __exit__(self, unused_type, unused_value, unused_traceback):
+    assert _control_ctx()[-1] is self
+    _control_ctx().pop()
+
+
+class NullCtx(object):
+  """Helper substitute for contextlib.nullcontext."""
+
+  def __enter__(self):
+    pass
+
+  def __exit__(self, unused_type, unused_value, unused_traceback):
+    pass
+
+
+def _default_control_status_ctx():
+  return ControlStatusCtx(status=Status.UNSPECIFIED)
+
+
+INSPECT_SOURCE_SUPPORTED = True
+try:
+  inspect.getsource(ag_logging.log)
+except OSError:
+  INSPECT_SOURCE_SUPPORTED = False
+  ag_logging.warning(
+      'AutoGraph is not available in this environment: functions lack code'
+      ' information. This is typical of some environments like the interactive'
+      ' Python shell. See'
+      ' https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/autograph/g3doc/reference/limitations.md#access-to-source-code'
+      ' for more information.')

tensorflow/python/autograph/core/config.py

@@ -0,0 +1,62 @@
+# Copyright 2016 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.
+# ==============================================================================
+"""Global configuration."""
+
+from tensorflow.python.autograph.core import config_lib
+
+Action = config_lib.Action
+Convert = config_lib.Convert
+DoNotConvert = config_lib.DoNotConvert
+
+
+# This list is evaluated in order and stops at the first rule that tests True
+# for a definitely_convert of definitely_bypass call.
+CONVERSION_RULES = (
+    # Known packages
+    Convert('tensorflow.python.training.experimental'),
+
+    # Builtin modules
+    DoNotConvert('collections'),
+    DoNotConvert('copy'),
+    DoNotConvert('cProfile'),
+    DoNotConvert('inspect'),
+    DoNotConvert('ipdb'),
+    DoNotConvert('linecache'),
+    DoNotConvert('mock'),
+    DoNotConvert('pathlib'),
+    DoNotConvert('pdb'),
+    DoNotConvert('posixpath'),
+    DoNotConvert('pstats'),
+    DoNotConvert('re'),
+    DoNotConvert('threading'),
+    DoNotConvert('urllib'),
+
+    # Known libraries
+    DoNotConvert('matplotlib'),
+    DoNotConvert('numpy'),
+    DoNotConvert('pandas'),
+    DoNotConvert('tensorflow'),
+    DoNotConvert('PIL'),
+    DoNotConvert('absl.logging'),
+
+    # TODO(b/133417201): Remove.
+    DoNotConvert('tensorflow_probability'),
+
+    # TODO(b/133842282): Remove.
+    DoNotConvert('tensorflow_datasets.core'),
+
+    DoNotConvert('keras'),
+    DoNotConvert('tf_keras'),
+)

tensorflow/python/autograph/core/config_lib.py

@@ -0,0 +1,61 @@
+# Copyright 2016 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.
+# ==============================================================================
+"""Global configuration support."""
+
+import enum
+
+
+# TODO(mdan): For better performance, allow each rule to take a set names.
+
+
+class Rule(object):
+  """Base class for conversion rules."""
+
+  def __init__(self, module_prefix):
+    self._prefix = module_prefix
+
+  def matches(self, module_name):
+    return (module_name.startswith(self._prefix + '.') or
+            module_name == self._prefix)
+
+
+class Action(enum.Enum):
+  NONE = 0
+  CONVERT = 1
+  DO_NOT_CONVERT = 2
+
+
+class DoNotConvert(Rule):
+  """Indicates that this module should be not converted."""
+
+  def __str__(self):
+    return 'DoNotConvert rule for {}'.format(self._prefix)
+
+  def get_action(self, module):
+    if self.matches(module.__name__):
+      return Action.DO_NOT_CONVERT
+    return Action.NONE
+
+
+class Convert(Rule):
+  """Indicates that this module should be converted."""
+
+  def __str__(self):
+    return 'Convert rule for {}'.format(self._prefix)
+
+  def get_action(self, module):
+    if self.matches(module.__name__):
+      return Action.CONVERT
+    return Action.NONE

tensorflow/python/autograph/core/converter.py

@@ -0,0 +1,316 @@
+# Copyright 2016 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.
+# ==============================================================================
+"""Converter construction support.
+
+This module contains a base class for all converters, as well as supporting
+structures. These structures are referred to as contexts.
+
+The class hierarchy is as follows:
+
+    <your converter>
+      [extends] converter.Base
+        [extends] transformer.Base
+            [extends] gast.nodeTransformer
+          [uses] transformer.SourceInfo
+        [uses] converter.EntityContext
+          [uses] converter.ProgramContext
+          [uses] transformer.SourceInfo
+
+converter.Base is a specialization of transformer.Base for AutoGraph. It's a
+very lightweight subclass that adds a `ctx` attribute holding the corresponding
+EntityContext object (see below). Note that converters are not reusable, and
+`visit` will raise an error if called more than once.
+
+converter.EntityContext contains mutable state associated with an entity that
+the converter processes.
+
+converter.ProgramContext contains mutable state across related entities. For
+example, when converting several functions that call one another, the
+ProgramContext should be shared across these entities.
+
+Below is the overall flow at conversion:
+
+    program_ctx = ProgramContext(<entities to convert>, <global settings>, ...)
+    while <program_ctx has more entities to convert>:
+      entity, source_info = <get next entity from program_ctx>
+      entity_ctx = EntityContext(program_ctx, source_info)
+      for <each ConverterClass>:
+        converter = ConverterClass(entity_ctx)
+
+        # May update entity_ctx and program_ctx
+        entity = converter.visit(entity)
+
+      <add entity's dependencies to program_ctx>
+
+Note that pyct contains a small number of transformers used for static analysis.
+These implement transformer.Base, rather than converter.Base, to avoid a
+dependency on AutoGraph.
+"""
+
+import enum
+
+from tensorflow.python.autograph.pyct import anno
+from tensorflow.python.autograph.pyct import ast_util
+from tensorflow.python.autograph.pyct import parser
+from tensorflow.python.autograph.pyct import templates
+from tensorflow.python.autograph.pyct import transformer
+from tensorflow.python.util.tf_export import tf_export
+
+# TODO(mdan): These contexts can be refactored into first class objects.
+# For example, we could define Program and Entity abstractions that hold on
+# to the actual entity and have conversion methods.
+
+# TODO(mdan): Add a test specific to this converter.
+
+
+@tf_export('autograph.experimental.Feature')
+class Feature(enum.Enum):
+  """This enumeration represents optional conversion options.
+
+  These conversion options are experimental. They are subject to change without
+  notice and offer no guarantees.
+
+  _Example Usage_
+
+  ```python
+  optionals= tf.autograph.experimental.Feature.EQUALITY_OPERATORS
+  @tf.function(experimental_autograph_options=optionals)
+  def f(i):
+    if i == 0:  # EQUALITY_OPERATORS allows the use of == here.
+      tf.print('i is zero')
+  ```
+
+  Attributes:
+    ALL: Enable all features.
+    AUTO_CONTROL_DEPS: Insert of control dependencies in the generated code.
+    ASSERT_STATEMENTS: Convert Tensor-dependent assert statements to tf.Assert.
+    BUILTIN_FUNCTIONS: Convert builtin functions applied to Tensors to
+      their TF counterparts.
+    EQUALITY_OPERATORS: Whether to convert the equality operator ('==') to
+      tf.math.equal.
+    LISTS: Convert list idioms, like initializers, slices, append, etc.
+    NAME_SCOPES: Insert name scopes that name ops according to context, like the
+      function they were defined in.
+  """
+
+  ALL = 'ALL'
+
+  AUTO_CONTROL_DEPS = 'AUTO_CONTROL_DEPS'
+  ASSERT_STATEMENTS = 'ASSERT_STATEMENTS'
+  BUILTIN_FUNCTIONS = 'BUILTIN_FUNCTIONS'
+  EQUALITY_OPERATORS = 'EQUALITY_OPERATORS'
+  LISTS = 'LISTS'
+  NAME_SCOPES = 'NAME_SCOPES'
+
+  @classmethod
+  def all(cls):
+    """Returns a tuple that enables all options."""
+    return tuple(cls.__members__.values())
+
+  @classmethod
+  def all_but(cls, exclude):
+    """Returns a tuple that enables all but the excluded options."""
+    if not isinstance(exclude, (list, tuple, set)):
+      exclude = (exclude,)
+    return tuple(set(cls.all()) - set(exclude) - {cls.ALL})
+
+
+STANDARD_OPTIONS = None  # Forward definition.
+
+
+class ConversionOptions(object):
+  """Immutable container for global conversion flags.
+
+  Attributes:
+    recursive: bool, whether to recursively convert any user functions or
+      classes that the converted function may use.
+    user_requested: bool, whether the conversion was explicitly requested by
+      the user, as opposed to being performed as a result of other logic. This
+      value always auto-resets to False in child conversions.
+    optional_features: Union[Feature, Set[Feature]], controls the use of
+      optional features in the conversion process. See Feature for available
+      options.
+  """
+
+  def __init__(self,
+               recursive=False,
+               user_requested=False,
+               internal_convert_user_code=True,
+               optional_features=Feature.ALL):
+    self.recursive = recursive
+    self.user_requested = user_requested
+    # TODO(mdan): Rename to conversion_recursion_depth?
+    self.internal_convert_user_code = internal_convert_user_code
+
+    if optional_features is None:
+      optional_features = ()
+    elif isinstance(optional_features, Feature):
+      optional_features = (optional_features,)
+    optional_features = frozenset(optional_features)
+    self.optional_features = optional_features
+
+  def as_tuple(self):
+    return (self.recursive, self.user_requested,
+            self.internal_convert_user_code, self.optional_features)
+
+  def __hash__(self):
+    return hash(self.as_tuple())
+
+  def __eq__(self, other):
+    assert isinstance(other, ConversionOptions)
+    return self.as_tuple() == other.as_tuple()
+
+  def __str__(self):
+    return 'ConversionOptions[{}]'
+
+  def uses(self, feature):
+    return (Feature.ALL in self.optional_features or
+            feature in self.optional_features)
+
+  def call_options(self):
+    """Returns the corresponding options to be used for recursive conversion."""
+    return ConversionOptions(
+        recursive=self.recursive,
+        user_requested=False,
+        internal_convert_user_code=self.recursive,
+        optional_features=self.optional_features)
+
+  def to_ast(self):
+    """Returns a representation of this object as an AST node.
+
+    The AST node encodes a constructor that would create an object with the
+    same contents.
+
+    Returns:
+      ast.Node
+    """
+    if self == STANDARD_OPTIONS:
+      return parser.parse_expression('ag__.STD')
+
+    template = """
+      ag__.ConversionOptions(
+          recursive=recursive_val,
+          user_requested=user_requested_val,
+          optional_features=optional_features_val,
+          internal_convert_user_code=internal_convert_user_code_val)
+    """
+
+    def list_of_features(values):
+      return parser.parse_expression('({})'.format(', '.join(
+          'ag__.{}'.format(str(v)) for v in values)))
+
+    expr_ast = templates.replace(
+        template,
+        recursive_val=parser.parse_expression(str(self.recursive)),
+        user_requested_val=parser.parse_expression(str(self.user_requested)),
+        internal_convert_user_code_val=parser.parse_expression(
+            str(self.internal_convert_user_code)),
+        optional_features_val=list_of_features(self.optional_features))
+    return expr_ast[0].value
+
+
+STANDARD_OPTIONS = ConversionOptions(
+    recursive=True,
+    user_requested=False,
+    internal_convert_user_code=True,
+    optional_features=None)
+
+
+class ProgramContext(object):
+  """ProgramContext keeps track of converting function hierarchies.
+
+  Attributes:
+    options: ConversionOptions
+    autograph_module: Deprecated. Do not use.
+  """
+
+  def __init__(self, options, autograph_module=None):
+    self.options = options
+    self.autograph_module = autograph_module
+
+
+class Base(transformer.Base):
+  """All converters should inherit from this class.
+
+  Attributes:
+    ctx: EntityContext
+  """
+
+  def __init__(self, ctx):
+    super(Base, self).__init__(ctx)
+
+    self._used = False
+    self._ast_depth = 0
+
+  def get_definition_directive(self, node, directive, arg, default):
+    """Returns the unique directive argument for a symbol.
+
+    See lang/directives.py for details on directives.
+
+    Example:
+       # Given a directive in the code:
+       ag.foo_directive(bar, baz=1)
+
+       # One can write for an AST node Name(id='bar'):
+       get_definition_directive(node, ag.foo_directive, 'baz')
+
+    Args:
+      node: ast.AST, the node representing the symbol for which the directive
+        argument is needed.
+      directive: Callable[..., Any], the directive to search.
+      arg: str, the directive argument to return.
+      default: Any
+
+    Raises:
+      ValueError: if conflicting annotations have been found
+    """
+    defs = anno.getanno(node, anno.Static.ORIG_DEFINITIONS, ())
+    if not defs:
+      return default
+
+    arg_values_found = []
+    for def_ in defs:
+      if (directive in def_.directives and arg in def_.directives[directive]):
+        arg_values_found.append(def_.directives[directive][arg])
+
+    if not arg_values_found:
+      return default
+
+    if len(arg_values_found) == 1:
+      return arg_values_found[0]
+
+    # If multiple annotations reach the symbol, they must all match. If they do,
+    # return any of them.
+    first_value = arg_values_found[0]
+    for other_value in arg_values_found[1:]:
+      if not ast_util.matches(first_value, other_value):
+        qn = anno.getanno(node, anno.Basic.QN)
+        raise ValueError(
+            '%s has ambiguous annotations for %s(%s): %s, %s' %
+            (qn, directive.__name__, arg, parser.unparse(other_value).strip(),
+             parser.unparse(first_value).strip()))
+    return first_value
+
+  def visit(self, node):
+    if not self._ast_depth:
+      if self._used:
+        raise ValueError('converter objects cannot be reused')
+      self._used = True
+
+    self._ast_depth += 1
+    try:
+      return super(Base, self).visit(node)
+    finally:
+      self._ast_depth -= 1

tensorflow/python/autograph/core/converter_test.py

@@ -0,0 +1,152 @@
+# 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 converter module."""
+
+import types
+
+from tensorflow.python.autograph.core import converter
+from tensorflow.python.autograph.core import converter_testing
+from tensorflow.python.autograph.pyct import anno
+from tensorflow.python.autograph.pyct import loader
+from tensorflow.python.autograph.pyct import parser
+from tensorflow.python.autograph.pyct import templates
+from tensorflow.python.platform import test
+
+
+class TestConverter(converter.Base):
+  pass
+
+
+class ConversionOptionsTest(converter_testing.TestCase):
+
+  def test_to_ast(self):
+    opts = converter.ConversionOptions()
+    opts_ast = opts.to_ast()
+
+    template = '''
+    def f():
+      return opts_ast
+    '''
+    opts_packed = templates.replace(template, opts_ast=opts_ast)
+
+    reparsed, _, _ = loader.load_ast(opts_packed)
+    fake_ag = types.ModuleType('fake_ag')
+    fake_ag.ConversionOptions = converter.ConversionOptions
+    fake_ag.Feature = converter.Feature
+    reparsed.ag__ = fake_ag
+
+    reparsed_opts = reparsed.f()
+
+    self.assertEqual(opts.recursive, reparsed_opts.recursive)
+    self.assertEqual(opts.user_requested, False)
+    self.assertEqual(
+        opts.internal_convert_user_code,
+        reparsed_opts.internal_convert_user_code)
+    self.assertEqual(opts.optional_features, reparsed_opts.optional_features)
+
+
+class ConverterBaseTest(converter_testing.TestCase):
+
+  def test_get_definition_directive_basic(self):
+
+    directive_key = object
+
+    def f():
+      a = 1
+      return a
+
+    _, node, ctx = self.transform(f, (), include_ast=True)
+
+    symbol_a = node.body[1].value
+    defs, = anno.getanno(symbol_a, anno.Static.ORIG_DEFINITIONS)
+    defs.directives[directive_key] = {
+        'test_arg': parser.parse_expression('foo'),
+        'other_arg': parser.parse_expression('bar'),
+    }
+    c = TestConverter(ctx)
+    value = c.get_definition_directive(symbol_a, directive_key, 'test_arg',
+                                       None)
+    self.assertEqual(value.id, 'foo')
+
+  def test_get_definition_directive_default(self):
+
+    directive_key = object
+
+    def f():
+      a = 1
+      return a
+
+    _, node, ctx = self.transform(f, (), include_ast=True)
+
+    symbol_a = node.body[1].value
+    c = TestConverter(ctx)
+    value = c.get_definition_directive(symbol_a, directive_key, 'test_arg',
+                                       parser.parse_expression('default'))
+    self.assertEqual(value.id, 'default')
+
+  def test_get_definition_directive_multiple_consistent(self):
+
+    directive_key = object
+
+    def f():
+      a = 1
+      if a:
+        a = 2
+      return a
+
+    _, node, ctx = self.transform(f, (), include_ast=True)
+
+    symbol_a = node.body[2].value
+    defs = anno.getanno(symbol_a, anno.Static.ORIG_DEFINITIONS)
+    defs[0].directives[directive_key] = {
+        'test_arg': parser.parse_expression('foo'),
+        'other_arg': parser.parse_expression('bar'),
+    }
+    defs[1].directives[directive_key] = {
+        'test_arg': parser.parse_expression('foo'),
+        'other_arg': parser.parse_expression('baz'),
+    }
+    c = TestConverter(ctx)
+    value = c.get_definition_directive(symbol_a, directive_key, 'test_arg',
+                                       None)
+    self.assertEqual(value.id, 'foo')
+
+  def test_get_definition_directive_multiple_inconsistent(self):
+
+    directive_key = object
+
+    def f():
+      a = 1
+      if a:
+        a = 2
+      return a
+
+    _, node, ctx = self.transform(f, (), include_ast=True)
+
+    symbol_a = node.body[2].value
+    defs = anno.getanno(symbol_a, anno.Static.ORIG_DEFINITIONS)
+    defs[0].directives[directive_key] = {
+        'test_arg': parser.parse_expression('foo'),
+    }
+    defs[1].directives[directive_key] = {
+        'test_arg': parser.parse_expression('bar'),
+    }
+    c = TestConverter(ctx)
+    with self.assertRaises(ValueError):
+      c.get_definition_directive(symbol_a, directive_key, 'test_arg', None)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/autograph/core/converter_testing.py

@@ -0,0 +1,126 @@
+# 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.
+# ==============================================================================
+"""Base class for tests in this module."""
+
+import contextlib
+import inspect
+import io
+import sys
+import types
+
+from tensorflow.python.autograph.core import config
+from tensorflow.python.autograph.core import converter
+from tensorflow.python.autograph.impl import api
+from tensorflow.python.framework import ops
+from tensorflow.python.platform import test
+
+
+def allowlist(f):
+  """Helper that marks a callable as whtelitisted."""
+  if 'allowlisted_module_for_testing' not in sys.modules:
+    allowlisted_mod = types.ModuleType('allowlisted_module_for_testing')
+    sys.modules['allowlisted_module_for_testing'] = allowlisted_mod
+    config.CONVERSION_RULES = (
+        (config.DoNotConvert('allowlisted_module_for_testing'),) +
+        config.CONVERSION_RULES)
+
+  f.__module__ = 'allowlisted_module_for_testing'
+
+
+def is_inside_generated_code():
+  """Tests whether the caller is generated code. Implementation-specific."""
+  frame = inspect.currentframe()
+  try:
+    frame = frame.f_back
+
+    internal_stack_functions = ('converted_call', '_call_unconverted')
+    # Walk up the stack until we're out of the internal functions.
+    while (frame is not None and
+           frame.f_code.co_name in internal_stack_functions):
+      frame = frame.f_back
+    if frame is None:
+      return False
+
+    return 'ag__' in frame.f_locals
+  finally:
+    del frame
+
+
+class TestingTranspiler(api.PyToTF):
+  """Testing version that only applies given transformations."""
+
+  def __init__(self, converters, ag_overrides):
+    super(TestingTranspiler, self).__init__()
+    if isinstance(converters, (list, tuple)):
+      self._converters = converters
+    else:
+      self._converters = (converters,)
+    self.transformed_ast = None
+    self._ag_overrides = ag_overrides
+
+  def get_extra_locals(self):
+    retval = super(TestingTranspiler, self).get_extra_locals()
+    if self._ag_overrides:
+      modified_ag = types.ModuleType('fake_autograph')
+      modified_ag.__dict__.update(retval['ag__'].__dict__)
+      modified_ag.__dict__.update(self._ag_overrides)
+      retval['ag__'] = modified_ag
+    return retval
+
+  def transform_ast(self, node, ctx):
+    node = self.initial_analysis(node, ctx)
+
+    for c in self._converters:
+      node = c.transform(node, ctx)
+
+    self.transformed_ast = node
+    self.transform_ctx = ctx
+    return node
+
+
+class TestCase(test.TestCase):
+  """Base class for unit tests in this module. Contains relevant utilities."""
+
+  def setUp(self):
+    # AutoGraph tests must run in graph mode to properly test control flow.
+    self.graph = ops.Graph().as_default()
+    self.graph.__enter__()
+
+  def tearDown(self):
+    self.graph.__exit__(None, None, None)
+
+  @contextlib.contextmanager
+  def assertPrints(self, expected_result):
+    try:
+      out_capturer = io.StringIO()
+      sys.stdout = out_capturer
+      yield
+      self.assertEqual(out_capturer.getvalue(), expected_result)
+    finally:
+      sys.stdout = sys.__stdout__
+
+  def transform(
+      self, f, converter_module, include_ast=False, ag_overrides=None):
+    program_ctx = converter.ProgramContext(
+        options=converter.ConversionOptions(recursive=True),
+        autograph_module=api)
+
+    tr = TestingTranspiler(converter_module, ag_overrides)
+    transformed, _, _ = tr.transform_function(f, program_ctx)
+
+    if include_ast:
+      return transformed, tr.transformed_ast, tr.transform_ctx
+
+    return transformed

tensorflow/python/autograph/core/function_wrappers.py

@@ -0,0 +1,113 @@
+# 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.
+# ==============================================================================
+"""Support for wrapping converted functions bodies with auxiliary logic."""
+
+from tensorflow.python.autograph.core import ag_ctx
+from tensorflow.python.autograph.core import converter
+from tensorflow.python.autograph.operators import variables
+from tensorflow.python.framework import auto_control_deps
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_util
+from tensorflow.python.util import nest
+
+
+# TODO(mdan): Move this into operators - it represents a function definition.
+
+
+class FunctionScope(object):
+  """Context manager that wraps the body of a converted function.
+
+  This context manager handles various operations related to the scope of a
+  function:
+    * optional TF name scopes - these name scopes match the name of the
+        function, for easy visualization in tensorBoard;
+    * optional automatic control dependencies - this adds the same mechanism
+        for control dependencies that is used by `@tf.function`; it can be
+        optionally enabled when using `tf.autograph.to_graph`;
+    * tracking of autograph conversion state (whether it's enabled by the user,
+        conversion options;
+  """
+
+  def __init__(self, function_name, scope_name, options):
+    self.name = scope_name
+    self.options = options
+
+    if options.user_requested:
+      self.autograph_ctx = ag_ctx.ControlStatusCtx(ag_ctx.Status.ENABLED,
+                                                   options)
+    self.callopts = options.call_options()
+
+    use_name_scope = options.uses(converter.Feature.NAME_SCOPES)
+    self.use_name_scope = use_name_scope
+    if use_name_scope:
+      self.name_scope = ops.name_scope(self._sanitize(function_name))
+
+    use_auto_deps = self.options.uses(converter.Feature.AUTO_CONTROL_DEPS)
+    self.use_auto_deps = use_auto_deps
+    if use_auto_deps:
+      self.autodeps_scope = auto_control_deps.AutomaticControlDependencies()
+      self._return_value_marked = False
+
+  def _sanitize(self, name):
+    """See https://www.tensorflow.org/api_docs/python/tf/Graph#name_scope."""
+    # TensorFlow doesn't like leading underscores at the top level.
+    if name and name.startswith('_'):
+      name = 'fn' + name
+    return name
+
+  def __enter__(self):
+    if self.options.user_requested:
+      self.autograph_ctx.__enter__()
+    if self.use_name_scope:
+      self.name_scope.__enter__()
+    if self.use_auto_deps:
+      self.autodeps_scope.__enter__()
+    return self
+
+  def __exit__(self, exc_type, exc_val, exc_tb):
+    if self.options.user_requested:
+      self.autograph_ctx.__exit__(exc_type, exc_val, exc_tb)
+    if self.use_name_scope:
+      self.name_scope.__exit__(exc_type, exc_val, exc_tb)
+    if self.use_auto_deps:
+      self.autodeps_scope.__exit__(exc_type, exc_val, exc_tb)
+
+  def ret(self, value, did_return):
+    """Marks a value as returned from the function guarded by the scope."""
+    del did_return
+
+    if isinstance(value, variables.UndefinedReturnValue):
+      return None
+
+    if self.use_auto_deps:
+      self._return_value_marked = True
+      if value is None:
+        # We don't create dummy returns, to preserve Python semantics. The user
+        # is responsible for adding a return value to the top-level function.
+        return None
+
+      def _mark_return_if_tensor(t):
+        if tensor_util.is_tf_type(t):
+          return self.autodeps_scope.mark_as_return(t)
+        return t
+
+      value = nest.map_structure(_mark_return_if_tensor, value)
+    return value
+
+
+def with_function_scope(thunk, scope_name, options):
+  """Inline version of the FunctionScope context manager."""
+  with FunctionScope('lambda_', scope_name, options) as scope:
+    return thunk(scope)

tensorflow/python/autograph/core/function_wrappers_test.py

@@ -0,0 +1,57 @@
+# 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 function_wrappers module."""
+
+from tensorflow.python.autograph.core import converter
+from tensorflow.python.autograph.core import function_wrappers
+from tensorflow.python.eager import context
+from tensorflow.python.framework import constant_op
+from tensorflow.python.ops import variables
+from tensorflow.python.platform import test
+
+
+class FunctionWrappersTest(test.TestCase):
+
+  def test_name_scope(self):
+    if context.executing_eagerly():
+      self.skipTest('Tensor names are disabled in eager')
+
+    with function_wrappers.FunctionScope(
+        'test_name', None,
+        converter.ConversionOptions(
+            optional_features=converter.Feature.NAME_SCOPES)):
+      t = constant_op.constant(1)
+    self.assertIn('test_name', t.name)
+
+  def test_auto_control_deps(self):
+    v = variables.Variable(1)
+    with function_wrappers.FunctionScope(
+        '_', None,
+        converter.ConversionOptions(
+            optional_features=converter.Feature.AUTO_CONTROL_DEPS)) as scope:
+      v.assign(2)
+      op = scope.ret(constant_op.constant(1), True)
+    self.evaluate(op)
+    self.assertEqual(self.evaluate(v.read_value()), 2)
+
+  def test_all_disabled(self):
+    with function_wrappers.FunctionScope(None, None,
+                                         converter.STANDARD_OPTIONS):
+      t = constant_op.constant(1)
+    self.assertEqual(self.evaluate(t), 1)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/autograph/core/unsupported_features_checker.py

@@ -0,0 +1,57 @@
+# 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.
+# ==============================================================================
+"""Checkers for detecting unsupported Python features."""
+
+import gast
+
+from tensorflow.python.autograph.pyct import errors
+
+
+class UnsupportedFeaturesChecker(gast.NodeVisitor):
+  """Quick check for Python features we know we don't support.
+
+  Any features detected will cause AutoGraph to not compile a function.
+  """
+
+  def visit_Attribute(self, node):
+    if (node.attr is not None
+        and node.attr.startswith('__') and not node.attr.endswith('__')):
+      raise errors.UnsupportedLanguageElementError(
+          'mangled names are not yet supported')
+    self.generic_visit(node)
+
+  def visit_For(self, node):
+    if node.orelse:
+      raise errors.UnsupportedLanguageElementError(
+          'for/else statement not yet supported')
+    self.generic_visit(node)
+
+  def visit_While(self, node):
+    if node.orelse:
+      raise errors.UnsupportedLanguageElementError(
+          'while/else statement not yet supported')
+    self.generic_visit(node)
+
+  # These checks could potentially be replaced with inspect.isgeneratorfunction
+  # to avoid a getsource/parse/ast-walk round trip.
+  def visit_Yield(self, node):
+    raise errors.UnsupportedLanguageElementError('generators are not supported')
+
+  def visit_YieldFrom(self, node):
+    raise errors.UnsupportedLanguageElementError('generators are not supported')
+
+
+def verify(node):
+  UnsupportedFeaturesChecker().visit(node)

tensorflow/python/autograph/g3doc/pyct_tutorial.ipynb

@@ -0,0 +1,482 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "fWfkYsCgPvqR"
+      },
+      "source": [
+        "# Short intro to the SCT library of AutoGraph\n",
+        "\n",
+        "**Work in progress, use with care and expect changes.**\n",
+        "\n",
+        "The `pyct` module packages the source code transformation APIs used by AutoGraph.\n",
+        "\n",
+        "This tutorial is just a preview - there is no PIP package yet, and the API has not been finalized, although most of those shown here are quite stable.\n",
+        "\n",
+        "[Run in Colab](https://colab.research.google.com/github/tensorflow/tensorflow/blob/master/tensorflow/python/autograph/g3doc/pyct_tutorial.ipynb)\n",
+        "\n",
+        "Requires `tf-nightly`:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "wq1DRamRlqoB"
+      },
+      "outputs": [],
+      "source": [
+        "!pip install tf-nightly"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "r7Q78WIKe2cu"
+      },
+      "source": [
+        "### Writing a custom code translator\n",
+        "\n",
+        "[transformer.CodeGenerator](https://github.com/tensorflow/tensorflow/blob/40802bcdb5c8a4379da2145441f51051402bd29b/tensorflow/python/autograph/pyct/transformer.py#L480) is an AST visitor that outputs a string. This makes it useful in the final stage of translating Python to another language."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "HHaCMFOpuoVx"
+      },
+      "source": [
+        "Here's a toy C++ code generator written using a `transformer.CodeGenerator`, which is just a fancy subclass of [ast.NodeVisitor](https://docs.python.org/3/library/ast.html#ast.NodeVisitor):"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "PJlTIbJlurpm"
+      },
+      "outputs": [],
+      "source": [
+        "import gast\n",
+        "from tensorflow.python.autograph.pyct import transformer\n",
+        "\n",
+        "class BasicCppCodegen(transformer.CodeGenerator):\n",
+        "\n",
+        "  def visit_Name(self, node):\n",
+        "    self.emit(node.id)\n",
+        "\n",
+        "  def visit_arguments(self, node):\n",
+        "    self.visit(node.args[0])\n",
+        "    for arg in node.args[1:]:\n",
+        "      self.emit(', ')\n",
+        "      self.visit(arg)\n",
+        "\n",
+        "  def visit_FunctionDef(self, node):\n",
+        "    self.emit('void {}'.format(node.name))\n",
+        "    self.emit('(')\n",
+        "    self.visit(node.args)\n",
+        "    self.emit(') {\\n')\n",
+        "    self.visit_block(node.body)\n",
+        "    self.emit('\\n}')\n",
+        "\n",
+        "  def visit_Call(self, node):\n",
+        "    self.emit(node.func.id)\n",
+        "    self.emit('(')\n",
+        "    self.visit(node.args[0])\n",
+        "    for arg in node.args[1:]:\n",
+        "      self.emit(', ')\n",
+        "      self.visit(arg)\n",
+        "    self.emit(');')\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "chCX1A3rA9Pn"
+      },
+      "source": [
+        "Another helpful API is [transpiler.GenericTranspiler](https://github.com/tensorflow/tensorflow/blob/ee7172a929cb0c3d94a094fafc60bbaa175c085d/tensorflow/python/autograph/pyct/transpiler.py#L227) which takes care of parsing:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "LmwWewU1Bw0B"
+      },
+      "outputs": [],
+      "source": [
+        "import gast\n",
+        "from tensorflow.python.autograph.pyct import transpiler\n",
+        "\n",
+        "class PyToBasicCpp(transpiler.GenericTranspiler):\n",
+        "\n",
+        "  def transform_ast(self, node, ctx):\n",
+        "    codegen = BasicCppCodegen(ctx)\n",
+        "    codegen.visit(node)\n",
+        "    return codegen.code_buffer"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "nUhlScyOjlYM"
+      },
+      "source": [
+        "Try it on a simple function:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "ty9q853QvUqo"
+      },
+      "outputs": [],
+      "source": [
+        "def f(x, y):\n",
+        "  print(x, y)\n",
+        "\n",
+        "code, _ = PyToBasicCpp().transform(f, None)\n",
+        "print(code)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "rmRI9dG_ydE_"
+      },
+      "source": [
+        "### Helpful static analysis passes\n",
+        "\n",
+        "The `static_analysis` module contains various helper passes for dataflow analysis.\n",
+        "\n",
+        "All these passes annotate the AST. These annotations can be extracted using [anno.getanno](https://github.com/tensorflow/tensorflow/blob/40802bcdb5c8a4379da2145441f51051402bd29b/tensorflow/python/autograph/pyct/anno.py#L111). Most of them rely on the `qual_names` annotations, which just simplify the way more complex identifiers like `a.b.c` are accessed.\n",
+        "\n",
+        "The most useful is the activity analysis which just inventories symbols read, modified, etc.:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "GEJ30Wea4Xfy"
+      },
+      "outputs": [],
+      "source": [
+        "def get_node_and_ctx(f):\n",
+        "  node, source = parser.parse_entity(f, ())\n",
+        "  f_info = transformer.EntityInfo(\n",
+        "    name='f',\n",
+        "    source_code=source,\n",
+        "    source_file=None,\n",
+        "    future_features=(),\n",
+        "    namespace=None)\n",
+        "  ctx = transformer.Context(f_info, None, None)\n",
+        "  return node, ctx"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "BiwPJrDd0aAX"
+      },
+      "outputs": [],
+      "source": [
+        "from tensorflow.python.autograph.pyct import anno\n",
+        "from tensorflow.python.autograph.pyct import parser\n",
+        "from tensorflow.python.autograph.pyct import qual_names\n",
+        "from tensorflow.python.autograph.pyct.static_analysis import annos\n",
+        "from tensorflow.python.autograph.pyct.static_analysis import activity\n",
+        "\n",
+        "\n",
+        "def f(a):\n",
+        "  b = a + 1\n",
+        "  return b\n",
+        "\n",
+        "\n",
+        "node, ctx = get_node_and_ctx(f)\n",
+        "\n",
+        "node = qual_names.resolve(node)\n",
+        "node = activity.resolve(node, ctx)\n",
+        "\n",
+        "fn_scope = anno.getanno(node, annos.NodeAnno.BODY_SCOPE)  # Note: tag will be changed soon.\n",
+        "\n",
+        "\n",
+        "print('read:', fn_scope.read)\n",
+        "print('modified:', fn_scope.modified)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "w8dBRlKkFNIP"
+      },
+      "source": [
+        "Another useful utility is the control flow graph builder.\n",
+        "\n",
+        "Of course, a CFG that fully accounts for all effects is impractical to build in a late-bound language like Python without creating an almost fully-connected graph. However, one can be reasonably built if we ignore the potential for functions to raise arbitrary exceptions."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "KvLe9lWnFg7N"
+      },
+      "outputs": [],
+      "source": [
+        "from tensorflow.python.autograph.pyct import cfg\n",
+        "\n",
+        "\n",
+        "def f(a):\n",
+        "  if a > 0:\n",
+        "    return a\n",
+        "  b = -a\n",
+        "\n",
+        "node, ctx = get_node_and_ctx(f)\n",
+        "\n",
+        "node = qual_names.resolve(node)\n",
+        "cfgs = cfg.build(node)\n",
+        "cfgs[node]"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "Cro-jfPA2oxR"
+      },
+      "source": [
+        "Other useful analyses include liveness analysis. Note that these make simplifying assumptions, because in general the CFG of a Python program is a graph that's almost complete. The only robust assumption is that execution can't jump backwards."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "73dARy4_2oAI"
+      },
+      "outputs": [],
+      "source": [
+        "from tensorflow.python.autograph.pyct import anno\n",
+        "from tensorflow.python.autograph.pyct import cfg\n",
+        "from tensorflow.python.autograph.pyct import qual_names\n",
+        "from tensorflow.python.autograph.pyct.static_analysis import annos\n",
+        "from tensorflow.python.autograph.pyct.static_analysis import reaching_definitions\n",
+        "from tensorflow.python.autograph.pyct.static_analysis import reaching_fndefs\n",
+        "from tensorflow.python.autograph.pyct.static_analysis import liveness\n",
+        "\n",
+        "\n",
+        "def f(a):\n",
+        "  b = a + 1\n",
+        "  return b\n",
+        "\n",
+        "\n",
+        "node, ctx = get_node_and_ctx(f)\n",
+        "\n",
+        "node = qual_names.resolve(node)\n",
+        "cfgs = cfg.build(node)\n",
+        "node = activity.resolve(node, ctx)\n",
+        "node = reaching_definitions.resolve(node, ctx, cfgs)\n",
+        "node = reaching_fndefs.resolve(node, ctx, cfgs)\n",
+        "node = liveness.resolve(node, ctx, cfgs)\n",
+        "\n",
+        "print('live into `b = a + 1`:', anno.getanno(node.body[0], anno.Static.LIVE_VARS_IN))\n",
+        "print('live into `return b`:', anno.getanno(node.body[1], anno.Static.LIVE_VARS_IN))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "GKSaqLbKQI_v"
+      },
+      "source": [
+        "### Writing a custom Python-to-Python transpiler\n",
+        "\n",
+        "`transpiler.Py2Py` is a generic class for a Python [source-to-source compiler](https://en.wikipedia.org/wiki/Source-to-source_compiler). It operates on Python ASTs. Subclasses override its [transform_ast](https://github.com/tensorflow/tensorflow/blob/95ea3404528afcb1a74dd5f0946ea8d17beda28b/tensorflow/python/autograph/pyct/transpiler.py#L261) method.\n",
+        "\n",
+        "Unlike the `transformer` module, which have an AST as input/output, the `transpiler` APIs accept and return actual Python objects, handling the tasks associated with parsing, unparsing and loading of code."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "eicHoYlzRhnc"
+      },
+      "source": [
+        "Here's a transpiler that does nothing:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "edaG6dWEPvUI"
+      },
+      "outputs": [],
+      "source": [
+        "from tensorflow.python.autograph.pyct import transpiler\n",
+        "\n",
+        "\n",
+        "class NoopTranspiler(transpiler.PyToPy):\n",
+        "\n",
+        "  def get_caching_key(self, ctx):\n",
+        "    # You may return different caching keys if the transformation may generate\n",
+        "    # code versions.\n",
+        "    return 0\n",
+        "\n",
+        "  def get_extra_locals(self):\n",
+        "    # No locals needed for now; see below.\n",
+        "    return {}\n",
+        "\n",
+        "  def transform_ast(self, ast, transformer_context):\n",
+        "    return ast\n",
+        "\n",
+        "tr = NoopTranspiler()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "hKxmlWeQSQyN"
+      },
+      "source": [
+        "The main entry point is the [transform](https://github.com/tensorflow/tensorflow/blob/95ea3404528afcb1a74dd5f0946ea8d17beda28b/tensorflow/python/autograph/pyct/transpiler.py#L384) method returns the transformed version of the input."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "HXTIYsunSVr1"
+      },
+      "outputs": [],
+      "source": [
+        "def f(x, y):\n",
+        "  return x + y\n",
+        "\n",
+        "\n",
+        "new_f, module, source_map = tr.transform(f, None)\n",
+        "\n",
+        "new_f(1, 1)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "aKO42LBXw3SD"
+      },
+      "source": [
+        "### Adding new variables to the transformed code\n",
+        "\n",
+        "The transformed function has the same global and local variables as the original function. You can of course generate local imports to add any new references into the generated code, but an easier method is to use the `get_extra_locals` method:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "_Wl0n5I_1NJZ"
+      },
+      "outputs": [],
+      "source": [
+        "from tensorflow.python.autograph.pyct import parser\n",
+        "\n",
+        "\n",
+        "class HelloTranspiler(transpiler.PyToPy):\n",
+        "\n",
+        "  def get_caching_key(self, ctx):\n",
+        "    return 0\n",
+        "\n",
+        "  def get_extra_locals(self):\n",
+        "    return {'name': 'you'}\n",
+        "\n",
+        "  def transform_ast(self, ast, transformer_context):\n",
+        "    print_code = parser.parse('print(\"Hello\", name)')\n",
+        "    ast.body = [print_code] + ast.body\n",
+        "    return ast\n",
+        "\n",
+        "\n",
+        "def f(x, y):\n",
+        "  pass\n",
+        "\n",
+        "new_f, _, _ = HelloTranspiler().transform(f, None)\n",
+        "\n",
+        "_ = new_f(1, 1)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "JcMSHJXK6pO2"
+      },
+      "outputs": [],
+      "source": [
+        "import inspect\n",
+        "\n",
+        "print(inspect.getsource(new_f))"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "collapsed_sections": [],
+      "last_runtime": {
+        "build_target": "",
+        "kind": "local"
+      },
+      "name": "pyctr_tutorial.ipynb",
+      "provenance": [
+        {
+          "file_id": "1dT93XRkt7vUpVp7GZech8LB0u1OytKff",
+          "timestamp": 1586205976756
+        }
+      ]
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

tensorflow/python/autograph/g3doc/reference/_control_flow_tutorial.ipynb

@@ -0,0 +1,862 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "-vLwpT31YOJk"
+      },
+      "source": [
+        "TODO(b/138297412): This colab retains some useful code snippets and demonstrations that used to be in the tf.function/AutoGraph customization tutorial, and should be rolled into the existing docs as part of a broader markdown->colab conversion."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "otIdN1TS8N7S"
+      },
+      "outputs": [],
+      "source": [
+        "import tensorflow as tf"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "I0xDjO4SHLUD"
+      },
+      "source": [
+        "Define a helper function to demonstrate the kinds of errors you might encounter:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "D25apou9IOXa"
+      },
+      "outputs": [],
+      "source": [
+        "import traceback\n",
+        "import contextlib\n",
+        "\n",
+        "# Some helper code to demonstrate the kinds of errors you might encounter.\n",
+        "@contextlib.contextmanager\n",
+        "def assert_raises(error_class):\n",
+        "  try:\n",
+        "    yield\n",
+        "  except error_class as e:\n",
+        "    print('Caught expected exception \\n  {}:'.format(error_class))\n",
+        "    traceback.print_exc(limit=2)\n",
+        "  except Exception as e:\n",
+        "    raise e\n",
+        "  else:\n",
+        "    raise Exception('Expected {} to be raised but no error was raised!'.format(\n",
+        "        error_class))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "5f05Vr_YBUCz"
+      },
+      "source": [
+        "## Using AutoGraph\n",
+        "\n",
+        "The [autograph](https://www.tensorflow.org/guide/function) library is fully integrated with `tf.function`, and it will rewrite conditionals and loops which depend on Tensors to run dynamically in the graph.\n",
+        "\n",
+        "`tf.cond` and `tf.while_loop` continue to work with `tf.function`, but code with control flow is often easier to write and understand when written in imperative style."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "xgKmkrNTZSyz"
+      },
+      "source": [
+        "## AutoGraph: Conditionals\n",
+        "\n",
+        "AutoGraph will convert `if` statements into the equivalent `tf.cond` calls.\n",
+        "\n",
+        "This substitution is made if the condition is a Tensor. Otherwise, the conditional is executed during tracing."
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "20WlM9T2I9EV"
+      },
+      "source": [
+        "Here is a function that checks if the resulting graph uses `tf.cond`:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "E-7KllizZYsy"
+      },
+      "outputs": [],
+      "source": [
+        "def test_tf_cond(f, *args):\n",
+        "  g = f.get_concrete_function(*args).graph\n",
+        "  if any(node.name == 'cond' for node in g.as_graph_def().node):\n",
+        "    print(\"{}({}) uses tf.cond.\".format(\n",
+        "        f.__name__, ', '.join(map(str, args))))\n",
+        "  else:\n",
+        "    print(\"{}({}) executes normally.\".format(\n",
+        "        f.__name__, ', '.join(map(str, args))))\n",
+        "\n",
+        "  print(\"  result: \",f(*args).numpy())"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "DlqiutEEJHOe"
+      },
+      "source": [
+        "This substitution is made if the condition is a Tensor. Otherwise, the conditional is executed during tracing.\n",
+        "\n",
+        "Passing a python `True` executes the conditional normally:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "fCMywOXwJLIQ"
+      },
+      "outputs": [],
+      "source": [
+        "@tf.function\n",
+        "def dropout(x, training=True):\n",
+        "  if training:\n",
+        "    x = tf.nn.dropout(x, rate=0.5)\n",
+        "  return x"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "68D2RZ17JM8u"
+      },
+      "outputs": [],
+      "source": [
+        "test_tf_cond(dropout, tf.ones([10], dtype=tf.float32), True)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "WEz0QYucJPBa"
+      },
+      "source": [
+        "But passing a tensor replaces the python `if` with a `tf.cond`:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "o86paGR-Zadi"
+      },
+      "outputs": [],
+      "source": [
+        "test_tf_cond(dropout, tf.ones([10], dtype=tf.float32), tf.constant(True))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "5xFLfdApZh8q"
+      },
+      "source": [
+        "`tf.cond` has a number of subtleties.\n",
+        "\n",
+        "it works by tracing both sides of the conditional, and then choosing the appropriate branch at runtime, depending on the condition. Tracing both sides can result in unexpected execution of Python code."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "VTMoZEVaZiwk"
+      },
+      "outputs": [],
+      "source": [
+        "@tf.function\n",
+        "def f(x):\n",
+        "  if x > 0:\n",
+        "    x = x + 1.\n",
+        "    print(\"Tracing `then` branch\")\n",
+        "  else:\n",
+        "    x = x - 1.\n",
+        "    print(\"Tracing `else` branch\")\n",
+        "  return x"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "HqBVIZWb0Qzn"
+      },
+      "outputs": [],
+      "source": [
+        "f(-1.0).numpy()"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "BIMfbXlW0QdP"
+      },
+      "outputs": [],
+      "source": [
+        "f(1.0).numpy()"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "2nBnJ42v0Pvq"
+      },
+      "outputs": [],
+      "source": [
+        "f(tf.constant(1.0)).numpy()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "zyzzvtN5Jfpb"
+      },
+      "source": [
+        "It requires that if one branch creates a tensor used downstream, the other branch must also create that tensor."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "k_dxWHeFZlaQ"
+      },
+      "outputs": [],
+      "source": [
+        "@tf.function\n",
+        "def f():\n",
+        "  if tf.constant(True):\n",
+        "    x = tf.ones([3, 3])\n",
+        "  return x\n",
+        "\n",
+        "# Throws an error because both branches need to define `x`.\n",
+        "with assert_raises(ValueError):\n",
+        "  f()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "wP-LZP6cztnu"
+      },
+      "source": [
+        "If you want to be sure that a particular section of control flow is never converted by autograph, then explicitly convert the object to a python type so an error is raised instead: "
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "iG_VDavjzrzV"
+      },
+      "outputs": [],
+      "source": [
+        "@tf.function\n",
+        "def f(x, y):\n",
+        "  if bool(x):\n",
+        "    y = y + 1.\n",
+        "    print(\"Tracing `then` branch\")\n",
+        "  else:\n",
+        "    y = y - 1.\n",
+        "    print(\"Tracing `else` branch\")\n",
+        "  return y"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "kQ4CRP9T0rH2"
+      },
+      "outputs": [],
+      "source": [
+        "f(True, 0).numpy()"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "ww9tCzHy0rkv"
+      },
+      "outputs": [],
+      "source": [
+        "f(False, 0).numpy()"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "ppuV7iug0r7i"
+      },
+      "outputs": [],
+      "source": [
+        "with assert_raises(TypeError):\n",
+        "  f(tf.constant(True), 0.0)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "yho4J0a0ZkQS"
+      },
+      "source": [
+        "## AutoGraph and loops\n",
+        "\n",
+        "AutoGraph has a few simple rules for converting loops.\n",
+        "\n",
+        "- `for`: Convert if the iterable is a tensor\n",
+        "- `while`: Convert if the while condition depends on a tensor\n",
+        "\n",
+        "If a loop is converted, it will be dynamically unrolled with `tf.while_loop`, or in the special case of a `for x in tf.data.Dataset`, transformed into `tf.data.Dataset.reduce`.\n",
+        "\n",
+        "If a loop is _not_ converted, it will be statically unrolled "
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "OyzGNQAuZsky"
+      },
+      "outputs": [],
+      "source": [
+        "def test_dynamically_unrolled(f, *args):\n",
+        "  g = f.get_concrete_function(*args).graph\n",
+        "  if any(node.name == 'while' for node in g.as_graph_def().node):\n",
+        "    print(\"{}({}) uses tf.while_loop.\".format(\n",
+        "        f.__name__, ', '.join(map(str, args))))\n",
+        "  elif any(node.name == 'ReduceDataset' for node in g.as_graph_def().node):\n",
+        "    print(\"{}({}) uses tf.data.Dataset.reduce.\".format(\n",
+        "        f.__name__, ', '.join(map(str, args))))\n",
+        "  else:\n",
+        "    print(\"{}({}) gets unrolled.\".format(\n",
+        "        f.__name__, ', '.join(map(str, args))))\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "KFO1BSN9JkRP"
+      },
+      "source": [
+        "### For loops\n",
+        "\n",
+        "Here is a `tf.function` that demonstrates static unrolling:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "frecgTco_00V"
+      },
+      "outputs": [],
+      "source": [
+        "@tf.function\n",
+        "def for_in_range():\n",
+        "  x = 0\n",
+        "  for i in range(5):\n",
+        "    x += i\n",
+        "  return x\n",
+        "\n",
+        "test_dynamically_unrolled(for_in_range)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "PMdl0azc_5d4"
+      },
+      "outputs": [],
+      "source": [
+        "@tf.function\n",
+        "def for_in_tfrange():\n",
+        "  x = tf.constant(0, dtype=tf.int32)\n",
+        "  for i in tf.range(5):\n",
+        "    x += i\n",
+        "  return x\n",
+        "\n",
+        "test_dynamically_unrolled(for_in_tfrange)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "Q7tmncQTZt6_"
+      },
+      "outputs": [],
+      "source": [
+        "@tf.function\n",
+        "def for_in_tfdataset():\n",
+        "  x = tf.constant(0, dtype=tf.int64)\n",
+        "  for i in tf.data.Dataset.range(5):\n",
+        "    x += i\n",
+        "  return x\n",
+        "\n",
+        "test_dynamically_unrolled(for_in_tfdataset)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "eyPzDYiJAC8f"
+      },
+      "outputs": [],
+      "source": [
+        "@tf.function\n",
+        "def while_py_cond():\n",
+        "  x = 5\n",
+        "  while x > 0:\n",
+        "    x -= 1\n",
+        "  return x\n",
+        "\n",
+        "test_dynamically_unrolled(while_py_cond)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "l6s7aU-padY5"
+      },
+      "outputs": [],
+      "source": [
+        "@tf.function\n",
+        "def while_tf_cond():\n",
+        "  x = tf.constant(5)\n",
+        "  while x > 0:\n",
+        "    x -= 1\n",
+        "  return x\n",
+        "\n",
+        "test_dynamically_unrolled(while_tf_cond)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "dSr64Xn6ap-S"
+      },
+      "source": [
+        " If you have a `break` or early `return` clause that depends on a tensor, the top-level condition or iterable should also be a tensor.\n",
+        "\n",
+        "Compare the following examples:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "hG2Fe_OEAwpY"
+      },
+      "outputs": [],
+      "source": [
+        "@tf.function\n",
+        "def while_py_true_py_break(x):\n",
+        "  while True:  # py true\n",
+        "    if x == 0: # py break\n",
+        "      break\n",
+        "    x -= 1\n",
+        "  return x\n",
+        "\n",
+        "test_dynamically_unrolled(while_py_true_py_break, 5)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "Sr2cn5bY_E_9"
+      },
+      "outputs": [],
+      "source": [
+        "@tf.function\n",
+        "def buggy_while_py_true_tf_break(x):\n",
+        "  while True:   # py true\n",
+        "    if tf.equal(x, 0): # tf break\n",
+        "      break\n",
+        "    x -= 1\n",
+        "  return x\n",
+        "\n",
+        "with assert_raises(TypeError):\n",
+        "  test_dynamically_unrolled(buggy_while_py_true_tf_break, 5)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "Q-VirD-5avdZ"
+      },
+      "outputs": [],
+      "source": [
+        "@tf.function\n",
+        "def while_tf_true_tf_break(x):\n",
+        "  while tf.constant(True): # tf true\n",
+        "    if x == 0:  # py break\n",
+        "      break\n",
+        "    x -= 1\n",
+        "  return x\n",
+        "\n",
+        "test_dynamically_unrolled(while_tf_true_tf_break, 5)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "Upx5J0j8_Ldu"
+      },
+      "outputs": [],
+      "source": [
+        "@tf.function\n",
+        "def buggy_py_for_tf_break():\n",
+        "  x = 0\n",
+        "  for i in range(5):  # py for\n",
+        "    if tf.equal(i, 3): # tf break\n",
+        "      break\n",
+        "    x += i\n",
+        "  return x\n",
+        "\n",
+        "with assert_raises(TypeError):\n",
+        "  test_dynamically_unrolled(buggy_py_for_tf_break)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "GQHbodav_QMt"
+      },
+      "outputs": [],
+      "source": [
+        "@tf.function\n",
+        "def tf_for_py_break():\n",
+        "  x = 0\n",
+        "  for i in tf.range(5): # tf for\n",
+        "    if i == 3:  # py break\n",
+        "      break\n",
+        "    x += i\n",
+        "  return x\n",
+        "\n",
+        "test_dynamically_unrolled(tf_for_py_break)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "hyksHW9TCukR"
+      },
+      "source": [
+        "In order to accumulate results from a dynamically unrolled loop, you'll want to use `tf.TensorArray`.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "HJ3Vb3dXfefN"
+      },
+      "outputs": [],
+      "source": [
+        "batch_size = 2\n",
+        "seq_len = 3\n",
+        "feature_size = 4\n",
+        "\n",
+        "def rnn_step(inp, state):\n",
+        "  return inp + state\n",
+        "\n",
+        "@tf.function\n",
+        "def dynamic_rnn(rnn_step, input_data, initial_state):\n",
+        "  # [batch, time, features] -> [time, batch, features]\n",
+        "  input_data = tf.transpose(input_data, [1, 0, 2])\n",
+        "  max_seq_len = input_data.shape[0]\n",
+        "\n",
+        "  states = tf.TensorArray(tf.float32, size=max_seq_len)\n",
+        "  state = initial_state\n",
+        "  for i in tf.range(max_seq_len):\n",
+        "    state = rnn_step(input_data[i], state)\n",
+        "    states = states.write(i, state)\n",
+        "  return tf.transpose(states.stack(), [1, 0, 2])\n",
+        "  \n",
+        "dynamic_rnn(rnn_step,\n",
+        "            tf.random.uniform([batch_size, seq_len, feature_size]),\n",
+        "            tf.zeros([batch_size, feature_size]))"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "9gmLpHY-bkly"
+      },
+      "source": [
+        "### Gotcha's\n",
+        "\n",
+        "As with `tf.cond`, `tf.while_loop` also comes with a number of subtleties.\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "FJdfznhhKO7D"
+      },
+      "source": [
+        "#### Zero iterations\n",
+        "\n",
+        "Since a loop can execute 0 times, all tensors used downstream of the while_loop must be initialized above the loop.\n",
+        "\n",
+        "Here is an example of incorrect code:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "CocT5RHwblrQ"
+      },
+      "outputs": [],
+      "source": [
+        "@tf.function\n",
+        "def buggy_loop_var_uninitialized():\n",
+        "  for i in tf.range(3):\n",
+        "    x = i\n",
+        "  return x\n",
+        "\n",
+        "with assert_raises(ValueError):\n",
+        "  buggy_loop_var_uninitialized()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "ncr7tRZ1KWh9"
+      },
+      "source": [
+        "And the correct version:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "Wm7wIKXcCDGf"
+      },
+      "outputs": [],
+      "source": [
+        "@tf.function\n",
+        "def f():\n",
+        "  x = tf.constant(0)\n",
+        "  for i in tf.range(3):\n",
+        "    x = i\n",
+        "  return x\n",
+        "\n",
+        "f()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "CM7qXVY0KZHB"
+      },
+      "source": [
+        "#### Consistent shapes and types\n",
+        "\n",
+        "The shape/dtypes of all loop variables must stay consistent with each iteration.\n",
+        "\n",
+        "Here is an incorrect example that attempts to change a tensor's type:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "FSftc9cCbpAo"
+      },
+      "outputs": [],
+      "source": [
+        "@tf.function\n",
+        "def buggy_loop_type_changes():\n",
+        "  x = tf.constant(0, dtype=tf.float32)\n",
+        "  for i in tf.range(3): # Yields tensors of type tf.int32...\n",
+        "    x = i\n",
+        "  return x\n",
+        "\n",
+        "with assert_raises(TypeError):\n",
+        "  buggy_loop_type_changes()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "colab_type": "text",
+        "id": "M5l90NAHKsUM"
+      },
+      "source": [
+        "Here is an incorrect example that attempts to change a Tensor's shape while iterating:"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "kWF189prbuK0"
+      },
+      "outputs": [],
+      "source": [
+        "@tf.function\n",
+        "def buggy_concat():\n",
+        "  x = tf.ones([0, 10])\n",
+        "  for i in tf.range(5):\n",
+        "    x = tf.concat([x, tf.ones([1, 10])], axis=0)\n",
+        "  return x\n",
+        "\n",
+        "with assert_raises(ValueError):\n",
+        "  buggy_concat()"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 0,
+      "metadata": {
+        "colab": {},
+        "colab_type": "code",
+        "id": "miYnYcznCHeV"
+      },
+      "outputs": [],
+      "source": [
+        "@tf.function\n",
+        "def concat_with_padding():\n",
+        "  x = tf.zeros([5, 10])\n",
+        "  for i in tf.range(5):\n",
+        "    x = tf.concat([x[:i], tf.ones([1, 10]), tf.zeros([4-i, 10])], axis=0)\n",
+        "    x.set_shape([5, 10])\n",
+        "  return x\n",
+        "\n",
+        "concat_with_padding()\n"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "collapsed_sections": [],
+      "name": "performance.ipynb",
+      "private_outputs": true,
+      "provenance": [],
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

tensorflow/python/autograph/g3doc/reference/common_errors.md

@@ -0,0 +1,163 @@
+# AutoGraph reference
+
+[Index](index.md)
+
+## Common AutoGraph errors
+
+### "WARNING: AutoGraph could not transform `<name>`"
+
+This warning is output when AutoGraph could not convert a function, for an
+unexpected reason. The error message contains the reason why the function could
+not be converted, as well as guidance on how to proceed next.
+
+The exact error message may vary from version to version but in general, the
+cause of the failure appears somewhere in the text, for example as
+"Cause: could not get source code" or "Original error: could not get source
+code".
+
+Note: AutoGraph does not always output a warning. For example, constructors
+are silently called without conversion.
+
+When this warning is printed, the code returned by AutoGraph still executes, but
+the functions indicated in the warning will be executed as they are, without
+conversion. If the functions contain pure Python or graph code (for example,
+they have no Tensor-dependent control flow), then the code is likely to still
+run without error. However, if it contains any constructs that are only
+supported in AutoGraph, expect subsequent exceptions.
+
+Note: the warning is output to the [abseil](https://github.com/abseil/abseil-py)
+logger, with `WARNING` severity. To direct these warnings to `stdout`, use
+`tf.autograph.set_verbosity(0, True)`.
+
+### "Unable to locate the source code" or "Source not found" errors
+
+Newer versions of AutoGraph raise a `ConversionError`. Older versions print a
+warning. In both cases, a similar message about finding the source code is
+included.
+
+These errors are raised when AutoGraph is unable to find the source code of
+functions it needs to transform. See [Limitations](limitations.md) for more
+details.
+
+### "WARNING: Large unrolled loop detected"
+
+This warning is output when AutoGraph detects a `for` or `while` loop that
+creates TensorFlow ops and which has a large number of iterations and creates.
+
+This usually indicates a loop that was intended to run as a `tf.while_loop`, but
+instead runs as a Python loop.
+
+For example, a training loop might mistakenly iterate over a Python `range`,
+instead of `tf.range`:
+
+```
+num_steps = 10000
+step = tf.constant(0)
+for i in range(num_steps):
+  step += 1
+  train_step(model)
+```
+
+Another example is when using custom generators which AutoGraph does not
+support, even if they wrap over supported iterators like Datasets:
+
+```
+def my_iterator(ds):
+  for data in ds:
+    yield data
+
+# Custom iterators always dispatch to a Python for loop.
+for x in my_iterator(tf.data.Dataset.range(10)):
+  tf.print(x)
+```
+
+Note: This verification is only performed when `__debug__` is `True`.
+
+Note: the warning is output to the [abseil](https://github.com/abseil/abseil-py)
+logger, with `WARNING` severity. To direct these warnings to `stdout`, use
+`tf.autograph.set_verbosity(0, True)`.
+
+### "OperatorNotAllowedInGraphError: using a `tf.Tensor` as a Python `bool`"
+
+This exception is raised whenever a `tf.Tensor` is type-cast as a Python `bool`,
+in a context where eager execution is not active. The exception is only raised
+when graph execution is active, for example inside a `@tf.function` with
+AutoGraph turned off.
+
+**When AutoGraph is on**, it can be caused by: * placing a Tensor-dependent
+`break`, `continue` or `return` inside a Python loop (see example below) *
+attempting to use a `tf.Tensor` in a list comprehension, by iterating over it or
+using it in a condition
+
+A typical example of mixing Python and TF control flow in an incompatible way
+is:
+
+```
+for i in range(3):  # Python loop
+  if i > tf.constant(0):  # TF conditional
+    break  # raises OperatorNotAllowedInGraphError
+```
+
+The way these errors are typically fixed is by ensuring all control flow is
+TF control flow:
+
+```
+for i in tf.range(3):  # TF loop
+  if i > tf.constant(0):  # TF conditional
+    break  # works
+```
+
+**When AutoGraph is off**, it can be caused by using a `tf.Tensor` value as:
+
+  * the condition of an `if` or `while` statement: `if <tensor>:`
+  * the argument in a logical expression: `tensor and another_tensor`
+  * the argument to the `bool` built-in: `bool(tensor)`
+
+Note: These operations are allowed when executing eagerly.
+
+When encountering this error, make sure that the function is either decorated
+with `@tf.function`, or called from another function decorated in this way. Also
+look at the console and logging output for conversion warnings (see the section
+above).
+
+### "OperatorNotAllowedInGraphError: iterating over `tf.Tensor`"
+
+This exception is raised whenever you try to iterate over a `tf.Tensor`,
+in a context where eager execution is not active. The exception is only raised
+when graph execution is active, for example inside a `@tf.function` with
+AutoGraph turned off. It can be caused by using a `tf.Tensor` value as:
+
+  * the iterated of a `for` statement: `for i in tensor:`
+  * the argument to the `iter` built-in: `iter(tensor)`
+
+Note: These operations are allowed when executing eagerly.
+
+This exception is similar to the previous example, and has similar causes and
+remedies.
+
+### "InaccessibleTensorError: The tensor `<name>` is defined in another function or code block"
+
+This exception is common to code which attempts to obtain values calculated
+within a `tf.cond`, `tf.while_loop`, or another `@tf.function` without using
+functional style or through mutable collections. See
+[Capturing External Symbolic Tensors](https://www.tensorflow.org/guide/function#all_outputs_of_a_tffunction_must_be_return_values)
+and [Limitations](limitations.md) for more details.
+
+### "StagingError: in converted code"
+
+This exception is used by AutoGraph to wrap exceptions with custom constructors
+that it cannot re-raise with the original type. See
+[Error handling](error_handling.md) for more details. If your code uses custom
+exceptions, expect them to be wrapped by this exception.
+
+### "Unable to identify source code of lambda function"
+
+This error usually appears in the context of a conversion warning. It indicates
+that a lambda function could not be parsed (see [Limitations](limitations.md)).
+
+This type of error can usually be avoided by creating lambda functions in
+separate simple assignments, for example:
+
+```
+l = lambda <args>: <body>
+```

tensorflow/python/autograph/g3doc/reference/control_flow.md

@@ -0,0 +1,541 @@
+# AutoGraph reference
+
+[Index](index.md)
+
+## Control flow
+
+AutoGraph rewrites all control flow statements with specialized AutoGraph
+function calls. These function calls are capable of executing the corresponding
+control flow statement using Python semantics for effects outside the Python
+interpreter itself (see the [Introduction](intro.md)).
+
+### Dispatch rules
+
+Key Point: Only statements that are conditioned on, or iterate over, a
+TensorFlow object such as `tf.Tensor`, are converted into TensorFlow ops.
+
+As described in the [Introduction](intro.md), AutoGraph aims to preserve the
+semantics of valid Python code. If a control flow statement runs in graph
+execution without raising an error, then AutoGraph will also execute it as
+normal Python control flow. Statements which would normally raise an error, for
+example an `if` statement using a `bool` `Tensor` as condition, are converted to
+TensorFlow control flow ops.
+
+#### Analogy with compile-time constants and code optimization
+
+From the perspective of a TensorFlow graph, non-Tensor values, for example an
+integer or a NumPy array, are _constants_: they do not change value while the
+graph executes.
+
+For example, in the graph below, the condition is always `True` (it is
+invariant):
+
+```
+x = 1
+y = tf.cond(x > 0, lambda: 3 * x, lambda 5 * x)
+```
+
+That is equivalent to the code below:
+
+```
+x = 1
+y = 3 * x
+```
+
+In the example above, we've optimized away the conditional on a constant
+condition. The AutoGraph dispatch rules have the same effect: anything that is
+not a TensorFlow object is a compile-time constant for TensorFlow, and can be
+optimized away. For this reason, you can usually mix Python and TensorFlow
+computation, and it will transparently have the expected result even when only
+some computations are executed in the graph.
+
+<!-- TODO(mdan): This is actually a limitation (a very subtle one) -->
+Caution: The assumption of invariant code made above is not true if the
+TensorFlow graph had callbacks into the Python code. If you modify data
+from within a `tf.py_function`, then the code outside a `tf.py_function`
+will have unpredictable behavior if it depends on the same data.
+
+For example, the `tf.cond` that runs as part of the `if` statement below will
+miss the update made by `f`:
+
+```
+n = [10]
+def f():
+  n[0] = 20
+  return 0
+tf.py_function(f, (), (tf.int32,))
+if tf.equal(n[0], 10):
+  tf.print('n is 10')
+```
+
+```
+n is 10
+```
+
+### Compound symbols
+
+AutoGraph usually handles basic symbols:
+
+```
+if a < 0:
+  a = -a
+```
+
+```
+a = tf.cond(a < 0, lambda: -a, lambda: a)
+```
+
+But it can also handle complex symbols in many cases. For example, if we treat
+`a.b` as a symbol in the code below, then we can use it as if it were a basic
+symbol name:
+
+```
+if a.b < 0
+  a.b = -a.b
+```
+
+```
+a.b = tf.cond(a.b < 0, lambda: -a.b, lambda: a.b)
+```
+
+This is useful in methods, which can operate on properties of `self`, as well as
+working directly on more complex object structures or collections.
+
+Caution: There are certain [limitations](limitations.md) around using Python
+collections and object mutation. When in doubt, place the values you work
+with into local variables and operate on those.
+
+### Effects of the tracing process
+
+#### All Python code paths are executed during tracing
+
+When constructing a graph, TensorFlow _traces_ the code. The tracing of control
+flow requires visiting _every possible code path_ (usually once).
+
+Note: In rare cases, the runtime may decide to trace some code paths several
+times. For example, the condition of a `while` statement may be executed twice,
+first with a temporary graph, to determine whether it evaluates to a
+`tf.Tensor`, then if it is a `tf.Tensor`, it's executed a second time in the
+proper graph.
+
+In other words, tracing executes both branches of an if statement. Similarly,
+the body of loops is executed once (even if the loop would otherwise not iterate
+at all).
+
+This explains why inserting `print` statements in an `if` statement produces
+this output:
+
+```
+print('before if')
+if tf.constant(True):
+  print('true branch')
+else:
+  print('false branch')
+print('after if')
+```
+
+```
+before if
+true branch
+false branch
+after if
+```
+
+Note: Control flow that is not executed as a TensorFlow graph is not traced. Its
+body will execute as expected.
+
+Example of code that runs as regular Python code:
+
+```
+print('before if')
+if True:  # Condition not a Tensor, running normally
+  print('true branch')
+else:
+  print('false branch')
+print('after if')
+```
+
+```
+before if
+true branch
+after if
+```
+
+#### Python values modified in TensorFlow control flow become Tensors
+
+If a symbol is modified in a TensorFlow control flow statement, then it becomes
+a `tf.Tensor`, even if it started off as a Python primitive value.
+
+For example, the conditional below will run as a `tf.cond` (its condition is a
+`tf.Tensor`), which in turn will cause `i` to become a `tf.Tensor`.
+
+```
+i = 0
+if tf.greater(i, 0):
+  i = 1
+# i is now a Tensor
+```
+
+### `if` statements
+
+`if` statements whose condition is a `tf.Tensor` are executed as TensorFlow
+conditionals by converting them to `tf.cond`:
+
+```
+if tf.random.uniform(()) > 0.5:
+  x = 1
+else:
+  x = 2
+```
+
+`if` statements whose condition is not a `tf.Tensor` are executed as normal
+Python:
+
+```
+if np.random.uniform() > 0.5:
+  x = 1
+else:
+  x = 2
+```
+
+`if` statements executed as TensorFlow conditionals are subject to restrictions
+(see [limitations](limitations.md)). All symbols affected by the statement and
+used thereafter must be:
+
+ * of a data type understood by TensorFlow
+ * defined in both branches
+ * of consistent dtypes in both branches, for TensorFlow entities
+ * of consistent structure in both branches, for static collections (such as
+   lists or tuples)
+
+### `while` statements
+
+`while` statements whose condition is a `tf.Tensor` are executed as TensorFlow
+loops by converting them to `tf.while_loop`:
+
+```
+x = 0
+while tf.random.uniform(()) > 0.5:
+  x = x + 1
+```
+
+`while` statements whose condition is not a `tf.Tensor` are executed as normal
+Python:
+
+```
+x = 0
+while np.random.uniform() > 0.5:
+  x = x + 1
+```
+
+`while` statements executed as TensorFlow loops are subject to restrictions
+(see [limitations](limitations.md)). All symbols affected by the statement and
+used thereafter must be:
+
+ * of a data type understood by TensorFlow
+ * defined before the loop
+ * of consistent dtype at the beginning and the end of the loop,
+   for TensorFlow entities
+ * either of consistent shape at the beginning and the end of the loop,
+   for TensorFlow entities, or declared in `shape_invariants`
+ * of consistent structure  at the beginning and the end of the loop, for
+   static collections (such as lists or tuples)
+
+Caution: A `while` loop whose condition is a Python scalar will execute as
+normal Python. If you intended to run the loop as a TensorFlow loop, the loop
+will replicate its body in the graph (it is unrolled). To avoid that, make sure
+its condition is converted to a `tf.Tensor`, using for instance `tf.constant`.
+
+For example, the following loop is unrolled, even though the list contains
+`tf.Tensor` values, because the type of `l` is a Python `list`:
+
+```
+l = [tf.constant(1), tf.constant(2), tf.constant(3)]
+for i in l:
+  tf.print(i)  # This is unrolled - three `tf.print`s are built in the graph. 
+```
+
+If you wish for the loop to run as a TensorFlow loop, stack the loop:
+
+```
+l = [tf.constant(1), tf.constant(2), tf.constant(3)]
+for i in tf.stack(l):
+  tf.print(i)  # This runs as a TensorFlow loop.
+```
+
+<!-- TODO(mdan): List this under limitations -->
+
+Caution: A loop in which the type of the condition changes across iterations, in
+a way that would influence the way the loop is executed, is not allowed in
+AutoGraph.
+
+For example, the loop below will generate an error, because after the first
+iteration, `i` becomes a tf.Tensor:
+
+```
+i = 0
+while i < 10:  # `i < 10` is a Python bool - run as normal while loop
+  i = tf.constant(1)  # Error -- `i < 10` would now be a `tf.Tensor`
+```
+
+### `for` statements
+
+`for` statements that iterate over a `tf.Tensor` are executed as TensorFlow
+loops by converting them to a `tf.while_loop` which iterates over the first
+dimension (equivalent to NumPy):
+
+```
+for i in tf.constant(((1, 2), (3, 4))):
+  tf.print('iteration:', i)
+```
+
+```
+iteration: [1, 2]
+iteration: [3, 4]
+```
+
+Note: If possible, AutoGraph will also set the `maximum_iteration` parameter
+of the `tf.while_loop`.
+
+`for` statements that iterate over the output of a `tf.range` are executed as
+TensorFlow loops by converting them to a `tf.while_loop` which uses the
+arguments passed to the `tf.range`:
+
+```
+for i in tf.range(3):
+  tf.print('iteration:', i)
+```
+
+`for` statements that iterate over a `tf.data.Dataset` and which do not contain
+`break` or `return` statements are executed as TensorFlow loops by converting
+them to `tf.data.Dataset.reduce` ops:
+
+```
+for i in tf.data.Dataset.range(3):
+  tf.print('iteration:', i)
+```
+
+`for` statements that iterate over a _distributed dataset_ and which do not
+contain `break` or `return` statements are executed as TensorFlow loops by
+converting them to the dataset's `reduce` ops:
+
+```
+for i in tf.distribute.OneDeviceStrategy('cpu').experimental_distribute_dataset(
+    tf.data.Dataset.range(3)):
+  tf.print('iteration:', i)
+```
+
+`for` statements that iterate over a `tf.data.Dataset` and which contain
+`break` or `return` statements are executed as TensorFlow loops by converting
+them to a combination of `tf.data.Dataset.scan`, `tf.data.Dataset.take_while`
+and `tf.data.Dataset.reduce` ops:
+
+```
+for i in tf.data.Dataset.range(3):
+  tf.print('iteration:', i)
+  break
+```
+
+```
+iteration: 1
+```
+
+`for` statements that iterate over a `tf.data.Dataset` _iterator_ are executed
+as TensorFlow loops by converting them to a combination of `tf.while_loop`,
+and `tf.cond` ops:
+
+```
+for i in iter(tf.data.Dataset.range(3)):
+  tf.print('iteration:', i)
+```
+
+`for` statements that iterate over a type different from any of the above are
+executed as normal Python:
+
+```
+for i in [1, 2, 3]:
+  print('iteration:', i)
+```
+
+Caution: A `for` loop over a `list` or `tuple` of `tf.Tensor` is considered to
+iterate over a Python `list` (or respectively `tuple`), therefore will be
+executed as normal Python. If you intended to run it as a TensorFlow loop,
+use `tf.stack` or `tf.concat`.
+
+Caution: A `for` loop over a Python `range` will execute as normal Python.
+If you intended to run it as a TensorFlow loop, use `tf.range`.
+
+Note: AutoGraph may output a warning when it believes that you are unrolling
+a loop inefficiently. However, the warning thresholds are very conservative.
+The warning is only printed when
+[__debug__](https://docs.python.org/3/library/constants.html#__debug__) is
+`True`.
+
+Note: If `__debug__` is `True`, AutoGraph limits the number of iterations in
+normal Python loops to prevent infinite loops and raise an error if the limits
+are exceeded. However, the iteration limits are very large and may take a while
+to trigger an error.
+
+### `break` statements
+
+Code blocks in which `break` statements are used are rewritten with equivalent
+code that uses extra control booleans and conditionals. The control booleans are
+used directly in `while` loops. In the case of `for` loops, the AutoGraph
+corresponding operator accepts an `extra_test` argument which is similar to
+the conditional of a while loop, and which contains the control boolean.
+
+For example, the `while` loop below is rewritten as (showing the output of the
+`break` transformation only):
+
+```
+while i < 10:
+  if i > 3:
+    break
+  i += 1
+```
+
+```
+break_ = False
+while i < 10 and not break_:
+  if i > 3:
+    break_ = True
+    continue  # The continue statement is also rewritten in a subsequent pass
+  i += 1
+```
+
+Another example shows how the control boolean is used in the overload of a `for`
+loop (showing portions of the final output):
+
+```
+for i in range(10):
+  if i > 3:
+    break
+```
+
+```
+break_ = False
+...
+def extra_test(break_):
+  return ag__.not_(break_)
+# break_ becomes a loop variable.
+break_, = ag__.for_stmt(range(10), extra_test, ..., (break_,))
+```
+
+Mixing Tensor-dependent `break` and Python-dependent loops is disallowed:
+
+```
+@tf.function
+def buggy_while_py_true_tf_break(x):
+  while True:   # python conditional
+    if tf.equal(x, 0): # tensor break
+      break
+    x -= 1
+  return x
+
+# Raises OperatorNotAllowedInGraphError: using a `tf.Tensor` as a Python `bool` is not allowed
+# buggy_while_true_tf_break(5)
+```
+
+### `continue` statements
+
+Code blocks in which `continue` statements are used are rewritten with
+equivalent code that uses extra control booleans and conditionals, similar to
+how `break` is handled.
+
+For example, the `for` loop below is rewritten as (showing the output of the
+`continue` transformation only):
+
+```
+for i in range(10):
+  if i > 3:
+    continue
+```
+
+```
+for i in range(10):
+  continue_ = False
+  if i > 3:
+    continue_ = True
+  if not continue_:
+    i += 1
+```
+
+Notice that unlike `break`, `continue` statements are local to the loop and do
+not influence the number of iterations.
+
+### `return` statements
+
+`return` statements are also rewritten using control symbols, in a manner
+similar to how `break` is converted. In the case of `return` statements, an
+additional symbol keeps track of the return value.
+
+Depending on the structure of the code, the return value might be undefined
+in parts of the code (for example on code paths in which no return statement
+has executed). AutoGraph keeps track of this by using a special value.
+This special value is converted to `None` (the default return value) upon
+exiting the function.
+
+Caution: TensorFlow control flow does not support undefined values, and an
+undefined return value is no exception. Therefore, AutoGraph will raise an
+error for TensorFlow control flow in which the return value is not known for
+all code paths.
+
+For example, the following code raises an error because the return value would
+be undefined when the random number would be less than 0.5:
+
+```
+if tf.random.uniform(()) > 0.5:
+  return 1
+```
+
+```
+ValueError: A value must also be returned from the else branch.
+```
+
+An example of rewriting a `while` (showing the output of the `return`
+transformation only):
+
+```
+def f():
+  while i < 10:
+    if i > 3:
+      return 1
+    i += 1
+```
+
+```
+def f():
+  do_return = False
+  retval_ = ag__.UndefinedReturnValue()
+  while i < 10 and not do_return:
+    if i > 3:
+      do_return = True
+      retval_ = 1
+    if not do_return:
+      i += 1
+  return ag__.retval(retval_)  # Transforms any UndefinedReturnValue to None
+```
+
+Note: AutoGraph performs an additional code normalization in which an `if`
+statement with no `else` branch contains a `return` statement it is rewritten as
+an `if-else` statement in which the code that follows the statement is moved
+under the `else` branch.
+
+Example (showing the normalization only):
+
+```
+def f():
+  if i > 3:
+    return 1
+  i += 1
+```
+
+```
+def f():
+  if i > 3:
+    return 1
+  else:
+   i += 1
+```
+
+

tensorflow/python/autograph/g3doc/reference/debugging.md

@@ -0,0 +1,180 @@
+# AutoGraph reference
+
+[Index](index.md)
+
+## Debugging AutoGraph code
+
+The recommended way to debug AutoGraph code is to run it eagerly (see below).
+
+AutoGraph generates a new function, rather than directly executing the input
+function. Non-code elements, such as breakpoints, do not transfer to the
+generated code.
+
+You can step through the generated code and set breakpoints while debugging.
+The converted function is cached, and breakpoints should persist for the
+lifetime of the Python runtime.
+
+Note: The code generated by AutoGraph code is more complex than the input code,
+and is interspersed with AutoGraph boilerplate.
+
+Note: Python debugging can only be used to step through the code during graph
+construction time (or tracing time in the case of `tf.function`). To debug
+TensorFlow execution, use Eager execution.
+
+### Debugging `tf.function`: `tf.config.experimental_run_functions_eagerly`
+
+When using `@tf.function`, you can temporarily toggle graph execution by using
+`tf.config.experimental_run_functions_eagerly`. This will effectively run the
+annotated code eagerly, without transformation. Since AutoGraph has semantics
+consistent with Eager, it's an effective way to debug the code step-by-step.
+
+Note: AutoGraph is compatible with Eager, but the converse is not always
+true, so exercise care when making modifications to the code while debugging.
+
+Consider the following code:
+
+```
+@tf.function
+def f(a):
+  pdb.set_trace()
+  if a > 0:
+    tf.print(a, 'is positive')
+```
+
+Executing the line below will land the debugger in generated code, when the
+function is traced:
+
+```
+f(1)
+```
+
+```
+>l
+     10     def tf__f(a):
+     11       pdb.set_trace()
+---> 12       ag__.converted_call('print', tf, ag__.STD, (a,), None)
+     13
+     14       ...
+```
+
+Adding a call to `tf.config.experimental_run_functions_eagerly` before executing
+the function will land the debugger in the original code instead:
+
+```
+tf.config.run_functions_eagerly(True)
+f(1)
+```
+
+```
+>l
+      8 def f(a):
+      9   pdb.set_trace()
+---> 10   tf.print(a)
+     11   if a > 0:
+     12     tf.print('is positive')
+```
+
+### Using `print` and `tf.print`
+
+The `print` function is not converted by AutoGraph, and can be used to inspect
+the values of variables at graph construction time.
+
+Mixing `print` with `tf.print` can be confusing at first because they run at
+different stages. In general:
+
+ * all `print`s run when the TensorFlow graph is constructed
+ * all `tf.print`s run when the TensorFlow graph is executed
+
+#### Example: `print`
+
+To see the difference between `print` and `tf.print`.
+
+```
+@tf.function
+def f(a):
+  print(a)
+  if a > 0:
+    a = -a
+```
+
+When `a` is a `tf.Tensor` object, it is printed without an actual value:
+
+```
+f(tf.constant(1))
+```
+```
+Tensor("a:0", shape=(), dtype=int32)
+```
+
+Similarly, when `a` is just a Python value, it is printed directly:
+
+```
+f(1)
+```
+```
+1
+```
+
+#### Example: `print` followed by `tf.print`
+
+To see the difference between `print` and `tf.print`, let's run them together:
+
+```
+@tf.function
+def f(a):
+  print(a)
+  tf.print(a)
+```
+
+For non-`Tensor` values, they produce similar results:
+
+```
+f(1)
+```
+```
+1
+1
+```
+
+For Tensor values, only `tf.print` outputs the actual value:
+
+```
+f(tf.constant(1))
+```
+```
+Tensor("a:0", shape=(), dtype=int32)
+1
+```
+
+#### Example: `tf.print` followed by `print`
+
+Remember that, in general, *all `print`s run before all `tf.prints`*.
+What's more, since graphs are usually built once and executed multiple times,
+`print` usually runs just once when the function is first called.
+
+So in the example below, even though `tf.print` appears above `print`, it will
+run after it, because the graph is executed after it is built:
+
+```
+@tf.function
+def f(a):
+  tf.print('At graph execution:', a)
+  print('At graph construction:', a)
+```
+
+```
+f(tf.constant(1))
+```
+```
+At graph construction: Tensor("a:0", shape=(), dtype=int32)
+At graph execution: 1
+```
+
+Calling the function again will re-use the graph in this case:
+
+```
+f(tf.constant(1))
+```
+```
+At graph execution: 1
+```

tensorflow/python/autograph/g3doc/reference/error_handling.md

@@ -0,0 +1,217 @@
+# AutoGraph reference
+
+[Index](index.md)
+
+## Error handling
+
+When an exception occurs in code generated by AutoGraph, the error message
+is augmented with information about the location in the original code,
+before conversion.
+
+When an error occurs in a TensorFlow graph constructed using AutoGraph code,
+the stack trace which points to where the failing op was created is modified
+to point to the original code, before conversion.
+
+### Python execution errors
+
+Python execution (or tracing) exceptions that are raised in AutoGraph code are
+caught and re-raised with an extended error message that contains references
+to the original code.
+
+These functions are re-raised by `@tf.function`. If you use a `try/catch` the
+exception inside `tf.function`, you will obtain the original exception.
+
+The exception traceback still contains the entire call stack, including frames
+corresponding to generated code.
+
+AutoGraph tries to re-raise an exception of the same type as the original
+exception. This is usually possible for subclasses of
+`Exception` that do not define a custom `__init__`. For more complex
+exception types which define a custom constructor, AutoGraph raises a
+`StagingError` instead.
+
+Among the distinctive features of the re-raised exception:
+
+ * the exception traceback indicates the call stack of the exception, up to the
+   first @tf.function
+ * the error message includes references to the original code within
+   the `@tf.function`
+ * the references corresponding to converted code are marked with an
+   asterisk (`*`)
+ * the references corresponding to code which AutoGraph reached, but decided not
+   to convert, are marked with a double asterisk (`**`)
+ * the references corresponding to code that AutoGraph didn't reach at all have
+   no marking
+
+For example, the code below triggers an exception in the Python runtime, at
+graph construction time:
+
+```
+@tf.function
+def f():
+  tf.constant(1) + tf.constant(1.0)
+f()
+```
+
+An excerpt of the exception that is raised is shown below:
+
+```
+Traceback (most recent call last):
+  File "<ipython-input-10-1938a51c970d>", line 11, in <module>
+    f()
+  File "tensorflow/python/eager/def_function.py", line 417, in __call__
+    self._initialize(args, kwds, add_initializers_to=initializer_map)
+  ... more TensorFlow internal frames ...
+TypeError: in converted code:
+
+    <ipython-input-9-002fa22f79df>:8 f  *
+        tf.constant(1) + tf.constant(1.0)
+    tensorflow/python/ops/math_ops.py:900 binary_op_wrapper  **
+        return func(x, y, name=name)
+    ... more TensorFlow internal frames ...
+
+    TypeError: Input 'y' of 'AddV2' Op has type float32 that does not match type int32 of argument 'x'.
+
+```
+
+Note: the exact appearance of the various parts in the error message may change
+in the future.
+
+Let's look at the individual components of this exception.
+
+The traceback of the exception indicates the location until the call to
+`@tf.function`, including any frames internal to TensorFlow:
+
+```
+Traceback (most recent call last):
+  File "<ipython-input-10-1938a51c970d>", line 11, in <module>
+    f()
+  File "tensorflow/python/eager/def_function.py", line 417, in __call__
+    self._initialize(args, kwds, add_initializers_to=initializer_map)
+  File "tensorflow/python/eager/def_function.py", line 360, in _initialize
+    *args, **kwds))
+  File "tensorflow/python/eager/function.py", line 1688, in _get_concrete_function_internal_garbage_collected
+    graph_function, _, _ = self._maybe_define_function(args, kwargs)
+  File "tensorflow/python/eager/function.py", line 1992, in _maybe_define_function
+    graph_function = self._create_graph_function(args, kwargs)
+  File "tensorflow/python/eager/function.py", line 1878, in _create_graph_function
+    capture_by_value=self._capture_by_value),
+  File "tensorflow/python/framework/func_graph.py", line 791, in func_graph_from_py_func
+    func_outputs = python_func(*func_args, **func_kwargs)
+  File "tensorflow/python/eager/def_function.py", line 310, in wrapped_fn
+    return weak_wrapped_fn().__wrapped__(*args, **kwds)
+  File "tensorflow/python/framework/func_graph.py", line 781, in wrapper
+    raise e.ag_error_metadata.to_exception(type(e))
+```
+
+The exception message includes the location inside the converted function `f`:
+
+```
+TypeError: in converted code:
+
+    <ipython-input-9-002fa22f79df>:8 f  *
+        tf.constant(1) + tf.constant(1.0)
+    tensorflow/python/ops/math_ops.py:900 binary_op_wrapper
+        return func(x, y, name=name)
+    tensorflow/python/ops/math_ops.py:1198 _add_dispatch
+        return gen_math_ops.add_v2(x, y, name=name)
+    tensorflow/python/ops/gen_math_ops.py:549 add_v2
+        "AddV2", x=x, y=y, name=name)
+    tensorflow/python/framework/op_def_library.py:564 _apply_op_helper
+        inferred_from[input_arg.type_attr]))
+```
+
+Notice the frame corresponding to the call of `f`. The function is converted,
+which is being indicated by the asterisk `*` character displayed next to
+`f`:
+
+```
+    <ipython-input-9-002fa22f79df>:8 f  *
+        tf.constant(1) + tf.constant(1.0)
+```
+
+Lastly, the lower part includes the message that the exception originally
+reported:
+
+```
+    TypeError: Input 'y' of 'AddV2' Op has type float32 that does not match type int32 of argument 'x'.
+```
+
+Note: Typically, error messages raised by code internal to TensorFlow refers
+to arguments of the internal API that failed. Error messages raised by code
+internal to AutoGraph (that is, 'tensorflow/python/autograph') usually
+refer to symbols used in your code.
+
+### TensorFlow execution errors
+
+TensorFlow execution errors are displayed normally, but the portions of the
+error message which correspond to user code contain references to the original
+code.
+
+For example, the code below triggers an error in the TensorFlow runtime, at
+graph execution time:
+
+```
+@tf.function
+def my_function():
+  tf.Assert(tf.random.uniform(()) > 1.0, ['example error'])
+my_function()
+```
+
+An excerpt of the exception that is subsequently raised is shown below:
+
+```
+Traceback (most recent call last):
+  File "<ipython-input-16-af656fb445f0>", line 11, in <module>
+    my_function()
+  File "tensorflow/python/eager/def_function.py", line 435, in __call__
+    return self._concrete_stateful_fn._filtered_call(canon_args, canon_kwds)
+  File "tensorflow/python/eager/function.py", line 636, in _filtered_call
+    self.captured_inputs)
+  File "tensorflow/python/eager/function.py", line 734, in _call_flat
+    outputs = self._inference_function.call(ctx, args)
+  File "tensorflow/python/eager/function.py", line 460, in call
+    ctx=ctx)
+  File "tensorflow/python/eager/execute.py", line 68, in quick_execute
+    six.raise_from(core._status_to_exception(e.code, message), None)
+  File "<string>", line 3, in raise_from
+InvalidArgumentError:  assertion failed: [example error]
+    [[node Assert/Assert (defined at <ipython-input-16-af656fb445f0>:8) ]] [Op:__inference_my_function_79]
+```
+
+Notice the error message containing references to the location where the failing
+op was defined in the code (`<ipython-input-16-af656fb445f0>:8`):
+
+```
+InvalidArgumentError:  assertion failed: [example error]
+    [[node Assert/Assert (defined at <ipython-input-16-af656fb445f0>:8) ]] [Op:__inference_my_function_79]
+```
+
+### AutoGraph conversion exceptions
+
+Within `@tf.function`, when AutoGraph fails to convert a function, it displays
+a warning message and attempts to run the function without conversion.
+
+For example, the code below make a call to a Python
+[generator](https://wiki.python.org/moin/Generators) function, which is not
+supported by AutoGraph:
+
+```
+def example_generator():
+  yield 1
+
+@tf.function
+def f():
+  for i in example_generator():
+    print(i)
+```
+
+Calling `f()` will still run the code. AutoGraph will convert the function `f`,
+but skips the function `example_generator`. In addition, AutoGraph prints a
+warning to the console indicating that the function is called without being
+converted.
+
+```
+WARNING: Entity <function example_generator at 0x7f951b67f158> appears to be
+a generator function. It will not be converted by AutoGraph.
+```

tensorflow/python/autograph/g3doc/reference/functions.md

@@ -0,0 +1,67 @@
+# AutoGraph reference
+
+[Index](index.md)
+
+## Functions and function calls
+
+Typically, AutoGraph converts one function at a time. If a function calls other
+functions, the called function will be converted recursively, as described
+below.
+
+### Function calls
+
+AutoGraph rewrites all function calls with a special wrapper that may convert
+the called function at runtime.
+
+For example, the function call below:
+
+```
+f(x, y, z=1)
+```
+
+Is converted to code that schematically looks like this:
+
+```
+ag__.converted_call(f, ..., (x, y), {'z': 1}, ...)
+```
+
+All calls are rewritten, including calls to other types of callables, builtin
+functions, etc.
+
+If the originally called function is not converted, AutoGraph simply
+forwards the call to it, so that the wrapper is functionally equivalent with
+the original function call.
+
+If the originally called function is converted, then the conversion is performed
+first and the converted function is called instead.
+
+Note: a caching mechanism prevents the same function from being converted
+multiple times. This mechanism ensures that functions calls made with different
+[global or free variables](https://docs.python.org/3/reference/executionmodel.html#binding-of-names)
+are handled correctly.
+
+#### Function conversion rules
+
+The following types of functions are not converted:
+
+*   functions already converted
+*   functions defined in an allowlisted module (see autograph/core/config.py)
+*   non-Python functions (such as native bindings)
+*   `print`, `pdb.set_trace`, `ipdb.set_trace`
+*   most built-in functions (exceptions are listed in
+    autograph/operators/py_builtins.py)
+*   constructors
+*   functions without source code attached (prints a warning)(see
+    [limitations](limitations.md))
+*   generator functions (prints a warning)
+*   iterator protocol methods (`__next__`, `__iter__`)
+*   context manager methods (`__enter__`, `__exit__`)
+
+When AutoGraph encounters a function that it cannot convert outside of this
+list, it prints a warning.
+
+### Nested functions
+
+Functions nested inside a function converted by AutoGraph are converted
+at the same time as the function containing them. If the nested function is
+returned, a converted version of it is returned.

tensorflow/python/autograph/g3doc/reference/generated_code.md

@@ -0,0 +1,89 @@
+# AutoGraph reference
+
+[Index](index.md)
+
+## Generated code
+
+For each converted function, AutoGraph creates a new function. The
+loading mechanism is an implementation detail and may change, but the
+generated function is generally a regular
+[Python function](https://docs.python.org/3/reference/compound_stmts.html#function).
+This function is typically executed internally by `@tf.function` to construct a
+TensorFlow graph.
+
+### Transformations made to the generated code
+
+The generated code is a transformation of the input code. The transformations
+are listed below. Any other elements are left unchanged.
+
+Summary of transformations:
+
+ * function calls are replaced with a wrapper:
+   * `foo(args)` -> `ag__.converted_call(foo, args)`
+ * `if`, `while` and `for` statements are replaced with function calls:
+   * `if` -> `ag__.if_stmt`
+   * `while` -> `ag__.while_stmt`
+   * `for` -> `ag__.for_stmt`
+ * `break`, `return`, and `continue` statements are replaced with equivalent
+   `if` statements.
+ * `and`, `or` and `not` operators are replaced with function calls:
+   * `and` -> `ag__.and_`
+   * `or` -> `ag__.or_`
+   * `not` -> `ag__.not_`
+
+The functions replacing control flow statements are very similar in form with
+the corresponding control flow ops in TensorFlow.
+
+### AutoGraph generates normal Python code
+
+You can interact normally with the generated code. For example, you can use
+the `inspect.getsourcefile` and `inspect.getsource`:
+
+```
+def f(a):
+  ...
+
+converted_f = tf.autograph.to_graph(f)
+print(inspect.getsourcefile(converted_f))
+```
+```
+/tmp/tmpm562wlj7.py
+```
+
+When using `@tf.function`, you can repeat the same steps using the function's
+`python_function` attribute:
+
+```
+@tf.function
+def f(a):
+  ...
+
+converted_f = tf.autograph.to_graph(f.python_function)
+print(inspect.getsourcefile(converted_f))
+```
+```
+/tmp/tmpm562wlj7.py
+```
+
+`tf.autograph.to_code(f)` is a shortcut to obtain the generated code, and it's
+equivalent with calling `inspect.getsource(tf.autograph.to_graph(f))`.
+
+#### Recording diagnostic information: `tf.autograph.set_verbosity`
+
+AutoGraph can log additional debug information. This is mostly used for filing
+bugs, but can also be used to get an indication of whether a function is
+converted successfully or not.
+
+You can enable logging by calling `tf.autograph.set_verbosity(level)`. The
+`level` argument varies from 0 to 10:
+
+ * 0 - no logging
+ * 3 - includes the generated code
+ * 4 and above - extremely verbose logging
+
+Caution: The information being logged includes source code as well as
+data. Before sharing AutoGraph logs, make sure they don't contain any sensitive
+information.
+
+Alternatively, you can control the verbosity level using the environment
+variable `AUTOGRAPH_VERBOSITY`.

tensorflow/python/autograph/g3doc/reference/index.md

@@ -0,0 +1,23 @@
+# AutoGraph reference
+
+This reference document describes the semantics of AutoGraph transformations.
+
+In `@tf.function`, AutoGraph allows running Eager-style code as a TensorFlow
+graph.
+
+*   [Introduction](intro.md)
+*   [Interacting with the generated code](generated_code.md)
+*   [Debugging AutoGraph code](debugging.md)
+*   [Control flow](control_flow.md)
+*   [Functions and function calls](functions.md)
+*   [Error handling](error_handling.md)
+*   [Operator semantics](operators.md)
+*   [Limitations](limitations.md)
+*   [Common errors](common_errors.md)
+
+Also see:
+
+*   [AutoGraph blog post](https://medium.com/tensorflow/autograph-converts-python-into-tensorflow-graphs-b2a871f87ec7)
+*   [AutoGraph deep dive](https://www.youtube.com/watch?v=NIEgzljyDyI)
+*   [tf.function guide](https://www.tensorflow.org/guide/function)
+*   [Intro to TF graphs](https://www.tensorflow.org/guide/intro_to_graphs)

tensorflow/python/autograph/g3doc/reference/intro.md

@@ -0,0 +1,161 @@
+# AutoGraph reference
+
+[Index](index.md)
+
+## Introduction
+
+### Terminology
+
+Typically, AutoGraph operates by converting a function into a new function with
+new semantics.
+
+```
+def f(...):
+  ...
+
+converted_f = tf.autograph.to_graph(f)
+```
+
+In TensorFlow 2, AutoGraph is used to convert functions before they are being
+traced, when using the `tf.function` API. For example, the following code:
+
+```
+def f(...):
+  ...
+
+graph_f = tf.function(f)
+```
+
+is roughly equivalent to:
+
+```
+converted_f = tf.autograph.to_graph(f)
+graph_f = tf.function(autograph=False)(converted_f)
+```
+
+For the remainder of this document:
+
+ * **converted functions** are functions converted by AutoGraph
+ * **graph functions** are functions compiled with `tf.function`
+
+Graph functions are usually also converted, unless specified otherwise.
+
+### Safe behavior
+
+The semantics described below can be summarized as:
+
+ 1. code should either produce the same results as running it in Eager mode, or
+   fail with error.
+ 2. TF 1.* graph code should produce the same graph as running it directly in
+   graph mode.
+
+### Python semantics
+
+In general, AutoGraph does not change the semantics of error-free Python.
+In special circumstances, AutoGraph may also preserve the semantics of code that
+raises error as well. Semantics are preserved in the sense of the
+[as-if rule](https://en.wikipedia.org/wiki/As-if_rule).
+
+More specifically, code that would not be legal TensorFlow Graph code may become
+legal when converted.
+
+For example, consider the following function and the corresponding converted
+function:
+
+```
+def f(x):
+  if x > 0:
+    return x
+  return -x
+
+converted_f = tf.autograph.to_graph(f)
+```
+
+If `f` executes without error, then the `converted_f` has identical results:
+
+```
+>>> f(3)                               # Runs without error
+3
+>>> converted_f(3)                     # Identical result
+3
+```
+
+### TensorFlow Eager semantics
+
+If a function is called with `Tensor` arguments in Eager execution mode, it has
+certain semantics that are different from the Graph execution semantics.
+
+For example, the function below produces different results when executed in
+Eager and Graph modes:
+
+```
+def f(x):
+  if x > 0:
+    return x
+  return -x
+
+# Use tf.function to run code as a TensorFlow Graph.
+unconverted_graph_f = tf.function(f, autograph=False)
+graph_f = tf.function(f, autograph=True)
+```
+
+```
+>>> f(tf.constant(3))                      # Valid in Eager
+<tf.Tensor ... numpy=3>
+
+>>> unconverted_graph_f(tf.constant(3))    # Error - not legal Graph code
+TypeError: Using a `tf.Tensor` ...
+```
+
+AutoGraph transfers the Eager semantics to the graph function:
+
+```
+>>> graph_f(tf.constant(3))                # Valid in AutoGraph
+<tf.Tensor ... numpy=3>
+```
+
+### Subset of Eager
+
+AutoGraph currently supports only a subset of Eager. For example, list operation
+semantics are not supported:
+
+```
+def f(n):
+  l = []
+  for i in tf.range(n):
+    l.append(i)
+  return l
+
+converted_f = tf.autograph.to_graph(f)
+```
+
+```
+>>> f(3)                               # Valid in Eager
+[1, 2, 3]
+>>> converted_f(3)                     # Not valid in AutoGraph (by default)
+<<error>>
+```
+
+### Experimental features
+
+AutoGraph supports additional semantics which are not yet stable. This includes
+for example support for list semantics applied to `TensorArray`:
+
+```
+def f(n):
+  l = tf.TensorArray(dtype=tf.int32, size=0, dynamic_size=True)
+  for i in tf.range(n):
+    l.append(i)
+  return l
+
+converted_f = tf.autograph.to_graph(f)
+experimental_converted_f = tf.autograph.to_graph(
+    f, experimental_optional_features=tf.autograph.experimental.Feature.LISTS)
+```
+
+```
+>>> converted_f(3)                     # Not valid in AutoGraph (by default)
+<<error>>
+>>> experimental_converted_f(3)        # Valid
+<<TensorArray object>>
+```

tensorflow/python/autograph/g3doc/reference/limitations.md

@@ -0,0 +1,880 @@
+# AutoGraph reference
+
+[Index](index.md)
+
+## Limitations
+
+When AutoGraph is applied to normal Python code, you should expect no change
+in functionality.
+However, when applied to TensorFlow control flow (for example, an if statement
+with a `tf.Tensor` condition), there are certain limitations. This section
+describes these limitations and practices that will allow you to avoid them.
+
+Key Term: Python variables refer to Python symbols (or symbols for short) and
+should not be confused with TensorFlow variables.
+
+Key Term: A TensorFlow loop variable (or loop variable for short) refers to a
+value (typically a `tf.Tensor`) modified by a loop. See `tf.while_loop`.
+
+### Undefined and None values in TensorFlow
+
+TensorFlow does not support undefined or `None` values. All tensors must have
+a value.
+
+Example:
+
+```
+x = tf.cond(
+    tf.random.uniform(()) > 0.5,
+    lambda: tf.constant(1),
+    lambda: None)  # Error -- a Tensor cannot be None
+```
+
+The same restriction carries over in AutoGraph. If a variable is created inside
+control flow, and used after, then it must be defined before the control flow
+statement:
+
+```
+if tf.random.uniform(()) > 0.5:
+  x = tf.constant(1)
+else:
+  x = None
+tf.print(x)  # Error -- x may be None here
+```
+
+For this reason, AutoGraph forbids variables to be defined in only one branch
+of a TensorFlow conditional, if the variable is used afterwards:
+
+```
+del x
+if tf.random.uniform(()) > 0.5:
+  x = tf.constant(1)
+else:
+  pass
+tf.print(x)  # Error -- x may be undefined here
+```
+
+Note that if the variable is not used after the control flow statement, then it
+is considered local to the control flow block, and is not subject to these
+restrictions.
+
+```
+del x
+if tf.random.uniform(()) > 0.5:
+  x = tf.constant(1)  # Okay -- x does not need to be returned from the TF cond
+else:
+  pass
+```
+
+Similarly, variables must usually be defined before a TensorFlow loop.
+
+The most common example that is not allowed is a loop which initializes some
+accumulator variable in the first iteration:
+
+```
+del x
+for i in tf.range(100):  # Error -- x must be defined before the loop
+  if i == 0:
+    x = tf.constant(1)
+  else:
+    x = x + 1
+tf.print(x)
+```
+
+When the variable is only used inside the loop and does not depend on previous
+iterations, then it's ok to only be initialized inside the loop.
+
+```
+del x
+while tf.random.uniform(()) > 0.5:  # Okay -- x is not used after the loop
+  x = tf.constant(1)
+```
+
+* New in TF 2.4 *
+
+As long as it doesn't depend on previous iterations, the variable may also be
+used after the loop, however in that case the loop must execute at least one
+iteration, and will raise a runtime error otherwise.
+
+```
+del x
+for i in tf.range(10):  # Okay -- x does not depend on previous iterations
+  x = tf.constant(1)
+tf.print(x)
+```
+
+```
+del x
+while tf.constant(False):  # Error -- loop must initialize x!
+  x = tf.constant(1)
+tf.print(x)
+```
+
+Avoid these limitations by defining a default value before the control flow
+statement:
+
+```
+x = tf.constant()
+if tf.random.uniform(()) > 0.5:
+  x = tf.constant(1)
+tf.print(x)  # Okay -- x is either 0 or 1
+```
+
+Note: `None` values and undefined symbols are allowed in Eager control flow,
+because Eager execution uses Python control flow, rather than TensorFlow
+control flow ops.
+
+#### Special case: creating Tensors in a loop
+
+* New in TF 2.4 *
+
+A very common use-case is to run a training loop that creates some outputs:
+
+```
+for i in tf.range(num_steps):
+  outputs = train(next(data_iterator))
+```
+
+Often times these outputs can be nested structures of Tensors, which makes them
+impractical to initialize ahead of the loop.
+
+To help with this use-case, AutoGraph lets you run such loops, under certain
+conditions:
+
+ * outputs must be a Tensor, Python numeric, or a structure of these
+ * outputs must not depend on the value from a previous iteration; in other
+   words, the outputs may only appear to the left of an assignment operation
+ * the loop must run at least one iteration
+
+If the type of outputs is not recognized, then the usual
+"outputs must be defined before the loop" is raised at graph construction.
+
+AutoGraph also inserts a `tf.Assert` statement that raises a runtime error
+if the loop did not execute at least one iteration.
+
+### Indirect modifications and hidden side effects in TensorFlow control flow
+
+Key Point: We recommend using a functional programming style, immutable Python
+collections, TensorFlow ops and collections. Only TensorFlow objects should be
+used for side effects.
+
+#### AutoGraph analyzes code to detect modifications to Python objects
+
+Note: Modifications to TensorFlow objects, such as `tf.Variable`, are tracked
+using a different mechanism (automatic control dependencies) which does not
+rely on code analysis.
+
+One of the most important functions of AutoGraph is to rewrite Python control
+flow statements into equivalent TensorFlow ops. This process requires "wiring"
+variables covered by these control flow statements into the respective ops.
+
+The examples below use a `while` loop, but the same notions extend to all
+control flow such as `if` and `for` statements.
+
+In the example below, `x` needs to become a loop variable of the corresponding
+`tf.while_loop`:
+
+```
+while x > 0:
+  x = x - 1
+```
+```
+x = tf.while_loop(..., loop_vars=(x,)
+```
+
+TF control ops support only a limited set of types for loop variables. At the
+same time, the efficiency of TensorFlow graphs is influenced by the number of
+loop variables, so we don't want to create them unnecessarily. AutoGraph pulls
+symbols through loop variables only if necessary to minimize the number of
+loop variables.
+
+Note: If a symbol refers to a nested structure, such as a `dict` of `dict`s,
+the entire structure is mapped to multiple loop variables - TensorFlow
+automatically unpacks it.
+
+For example, the symbol 'y' below is not wired through the `tf.while_loop`'s
+`loop_vars` because it is not affected by the `while` loop:
+
+```
+y = 0
+while x > 0:
+  x = x - 1
+print(y)
+```
+```
+x = tf.while_loop(..., loop_vars=(x,)  # y does not need to be a loop variable
+```
+
+AutoGraph uses static analysis to determine which symbols are modified by the
+code, in order to transform them into control flow variables. Static analysis
+is generally performed on single functions - Python's dynamic nature limits its
+effectiveness across functions.
+
+#### Modifications of Python objects are not detected across functions
+
+Note: Modifications to TensorFlow objects, such as `tf.Variable`, are tracked
+using a different mechanism (automatic control dependencies). Modifications
+to `tf.Variable` objects are correctly handled even when called in other
+functions.
+
+Because static analysis is limited to single functions, modifications that are
+performed in other functions are not visible to AutoGraph:
+
+```
+def change_y():
+  global y
+  y = y + 1
+
+while x > 0:
+  change_y()  # Problem -- change made to y is not visible here!
+```
+
+This can be easily remedied using a functional programming style - writing
+functions that use argument for all their inputs and return values for all their
+outputs:
+
+```
+def change(y):
+  y = y + 1
+  return y
+
+while x > 0:
+  y = change(y)  # Okay -- y can now be properly tracked!
+```
+
+As noted before, this limitation does not apply to most TensorFlow objects,
+although it is still a good idea to use functional programming style for
+better code readability:
+
+```
+def change(y_var):
+  y_var.assign_add(1)
+
+y = tf.Variable(1)
+while x > 0:
+  change(y)  # This is still okay -- TensorFlow side effects are robust.
+```
+
+Keep in mind however that certain types like `tf.TensorArray` don't support
+side effects and must have their result assigned, otherwise they may raise an
+error:
+
+```
+def change(ta):
+  ta.write(0, 1)  # Incorrect use of TensorArray - will raise an error
+```
+
+In other words, `tf.TensorArray` must be handled using functional programming
+style:
+
+```
+def change(ta):
+  ta = ta.write(0, 1)  # Modifications create a new TensorArray efficiently.
+  return ta
+
+ta = tf.TensorArray(tf.int32, size=0, dynamic_size=True)
+while x > 0:
+  # TensorArray must be handled using functional programming style.
+  ta = change(ta)
+```
+
+#### Modifications of Python objects are not detected in methods
+
+A special case of hidden side effects are methods, which are commonly used
+to change the value of objects:
+
+```
+class MyClass(object):
+  def change(self):
+    self.y += 1
+
+c = MyClass()
+while x > 0:
+  c.change()  # Problem -- modification to c.y is not visible here!
+```
+
+This can be addressed in a number of ways.
+
+One possibility is to operate directly on the object properties:
+
+```
+c = MyClass()
+while x > 0:
+  c.y += 1  # Okay -- c.y can now be properly tracked!
+```
+
+Another possibility is to rely on immutable objects with value semantics. This
+may lead to many temporary objects when executing eagerly, but their number is
+greatly reduced in `@tf.function`:
+
+```
+class MyClass(collections.namedtuple('MyClass', ('y',))):
+  def change(self):
+    new_y = self.y + 1
+    return MyClass(new_y)
+
+c = MyClass()
+while x > 0:
+  c = c.change()  # Okay -- c is now a loop var.
+```
+
+It is also recommended to use a functional programming style with such immutable
+objects - that is, all arguments are inputs, all changes are return values:
+
+```
+def use_my_class(c: MyClass) -> MyClass:
+  new_c = c.change()
+  return new_c
+```
+
+Don't worry about creating a few extra objects - they are only used at trace
+time, and don't exist at graph execution.
+
+Note: TensorFlow control flow does not currently support arbitrary Python
+objects, but it does support basic collection objects such as `list`, `dict`,
+`tuple`, `namedtuple` and their subclasses. Design your objects as subclasses
+of [namedtuple](https://docs.python.org/3/library/collections.html#collections.namedtuple),
+or other types that [tf.nest](https://www.tensorflow.org/api_docs/python/tf/nest/map_structure)
+recognizes.
+
+#### Variables closed over by lambda functions
+
+### This limit will be deprecated after 2023-09-23
+
+AutoGraph assumes that variables that local functions close over may be used
+anywhere in the parent function, because in general it is possible to hide a
+function call in almost any Python statement. For this reason, these variables
+are accounted within TensorFlow loops.
+
+For example, the following code correctly captures `a` in the TensorFlow loop
+variables:
+
+```
+a = 0
+def f():
+  tf.print(a)
+for i in tf.range(3):
+  a = i
+f()  # Prints 2
+```
+
+A consequence is that these variables must be defined before the loop (see
+Undefined and None values above). So the following code will raise an error,
+even if the variable is never used after the loop:
+
+```
+def f():
+  tf.print(a)
+for i in tf.range(3):  # Error -- `a` must be defined before the loop.
+  a = i
+```
+
+However, lambda functions are handled differently, for reasons of backward
+compatibility. Lambda functions are assumed to be used in the statement where
+they are used, or at least in the same block.
+
+```
+a = 0
+foo(lambda: a)  # This lambda is not expected to be called anywhere else.
+for i in tf.range(3):  # Okay -- `a` is local to the loop.
+  a = i
+```
+
+Due to that reason, the following code will not work as expected for TensorFlow
+loops.
+
+```
+a = 0
+l = lambda: tf.print(a)
+for i in tf.range(3):
+  a = i  # `a` is considered local to the loop
+l()  # Prints 0!
+```
+
+Note that none of these restrictions only apply to TensorFlow loops; Python
+loops correctly handle closures in all cases.
+
+### Python collections in TensorFlow control flow
+
+Key Point: Use TensorFlow collection classes instead of Python collections.
+Python collections are okay to use when they represent a fixed structure (that
+is, `list`s don't change length, `dict`s don't add or remove keys).
+
+#### Modifying Python collections in TensorFlow control flow is not allowed
+
+One of the advantages of eager execution is that you may use the usual Python
+collections, like `list` or `dict` to hold `tf.Tensor` values. However, these
+are generally not compatible with TensorFlow control flow. Specialized
+collections like `tf.TensorArray` are required.
+
+Consider the following example:
+
+```
+def fn():
+  l = []
+
+  def loop_cond(i):
+    return i < 10
+
+  def loop_body(i):
+    i = i + 1
+    l.append(i)
+    return i,
+
+  tf.while_loop(
+      cond=loop_cond,
+      body=loop_body,
+      loop_vars=(0,))
+
+  return l
+```
+
+This code works in eager execution, which does not use the TensorFlow runtime
+for the `tf.while_loop`:
+
+```
+fn()
+```
+
+However, it does not work in graph execution, because TensorFlow uses special
+mechanisms to ensure the computations are correctly sequenced in the dataflow
+graph:
+
+```
+tf.function(fn)()  # Error -- illegal tensor capture!
+```
+
+The equivalent AutoGraph code raises the same error:
+
+```
+l = []
+for i in tf.range(10):
+  l.append(i)  # Error -- illegal tensor capture!
+```
+
+Instead, use the specialized `tf.TensorArray` class:
+
+```
+l = tf.TensorArray(tf.int32, size=0, dynamic_size=True)
+for i in tf.range(10):
+  l = l.write(l.size(), i)  # Okay
+```
+
+#### Python collections of fixed structure are allowed TensorFlow control flow
+
+An exception to the previous rule is made by Python collections that are static,
+that is, they don't grow in size for the duration of the computation.
+
+Caution: Use functional programming style when manipulating static collections.
+
+Examples:
+
+```
+static_list = [tf.constant(3)]
+while d.prop > 0:
+  static_list[0] -= 1  # Okay -- static_list does not change structure
+```
+```
+static_object = MyClass()
+static_object.field = tf.constant(3)
+while static_object.field > 0:
+  static_object.field -= 1  # Okay -- static_object does not change structure
+```
+```
+static_dict = {'field': tf.constant(3)}
+while static_dict['field'] > 0:
+  static_dict['field'] -= 1  # Okay -- static_dict does not change structure
+```
+
+However, remember to use functional programming style when these collections
+are used inside control flow.
+
+#### Python collections of fixed structure with dynamic index
+
+A more subtle error occurs when the collection is static, but is accessed in a
+dynamic way, that is with a key that is not constant.
+
+For example:
+
+```
+d = {'a': tf.constant(3)}
+for i in tf.range(10):
+  for key in d:
+    d[key] += i  # Problem -- accessing `dict` using non-constant key
+```
+
+The code above will raise an "illegal capture" error. To remedy it, write it in
+functional programming style:
+
+```
+d = {'a': tf.constant(3)}
+for i in tf.range(10):
+  d = {key: value + i for key, value in d.items()}  # Okay
+```
+
+### Shape and dtype consistency in TensorFlow control flow
+
+Unlike Python, TensorFlow has limited support for dynamic typing. This means
+that tensors must maintain consistent shapes and dtypes across control flow
+paths.
+
+Note: In general, these restrictions do not apply in control flow in Eager
+execution, because Eager execution uses Python control flow, rather than
+TensorFlow control flow ops.
+
+#### Mixing dynamic computations and static shapes
+
+Key Point: Use `.shape` on tensors of static shape, and `.shape.rank` on
+tensors of static rank; only use `tf.shape` and `tf.rank` when the shape or
+rank is dynamic.
+
+TensorFlow has optional static types and shapes: the shape of tensors may be
+static (e.g. `my_tensor.shape=(3, 3)` denotes a three by three matrix) or
+dynamic (e.g. `my_tensor.shape=(None, 3)`) denotes a matrix with a dynamic
+number of rows and three columns. When the shapes are dynamic, you can still
+query it at runtime by using the `tf.shape()` function.
+
+Note: `tf.shape` always returns a tensor.
+
+For static shapes, TensorFlow will perform additional shape verifications at
+graph construction time, that is, during tracing. These static shape
+verifications are useful because they work like a compiler for example, errors
+are caught early, before execution even begins.
+
+For example:
+
+```
+x = tf.constant([1, 2, 3])
+x[4]  # Tracing error! 4 is out of bounds.
+```
+
+To avoid tracing errors, you can add static shape verifications, which help
+make your code more robust:
+
+```
+if x.shape[0] > 4:
+  val = x[4]
+else:
+  val = some_default_value
+```
+
+In the snippet above, the code is protected against index-out-of-bounds
+errors. The code is also efficient because the verification `x.shape[0] > 4`
+will not be included in the graph.
+
+But what happens if you try to perform the index verifications using dynamic
+control flow? You might expect that the code works in the same way:
+
+```
+val = tf.cond(
+  x.shape[0] >= 4,
+  lambda: x[4],
+  lambda: some_default_value)
+```
+
+However, TensorFlow will not let you write code that could result in an error,
+even if that code appeared in a branch of a `tf.cond` statement that would
+never execute. Remember that the shape of `x` is `(3,)`, so TensorFlow performs
+static shape verification.
+
+This can lead to surprising behavior when using `tf.shape` on tensors with
+static shape in TensorFlow:
+
+```
+x = tf.constant((1, 2, 3))
+if tf.shape(x)[0] > 4:
+  val = x[4]  # Error at tracing: 4 is out of bounds!
+else:
+  val = some_default_value
+```
+
+Because `tf.shape` always evaluates to a Tensor, the `if` statement above is
+converted by AutoGraph into a `tf.cond`, which performs static shape
+verification of both branches.
+
+What if you need to write code which can handle both static and dynamic
+shapes? There are a few options in this case:
+
+A first option is to always work with dynamic shapes, for instance by
+using `input_signature` in `tf.function`. Many shape and shape-related checks
+are skipped when the shape is dynamic:
+
+```
+@tf.function(input_signature=(tf.TensorSpec(shape=(None,))))
+def f(x):  # x now has dynamic shape
+  if tf.shape(x)[0] >= 3:  # Builds a tf.cond
+    val = x[4]  # Okay, bounds checks are skipped when the shape is dynamic
+  else:
+    val = some_default_value
+```
+
+A second option is to first verify whether the shape is static or dynamic.
+This can be done at tracing time, allowing to use Python `if` to only trace
+the code that is suitable for the situation:
+
+```
+if x.shape[0] is None:  # Python bool, does not use tf.cond
+  # ... use x.shape here ...
+else:
+  # ... use tf.shape(x) here ...
+```
+
+#### Consistency of dtype
+
+The dtypes across all code paths must be consistent in conditionals and loops.
+
+For example, if a `tf.cond` (and correspondingly, an AutoGraph `if`) sets a
+tensor value conditionally, then that tensor must have the same shape and dtype
+in both branches of the conditional.
+
+Example of illegal dtype change in a conditional:
+
+```
+x = tf.cond(
+    tf.random.uniform(()) > 0.5,
+    lambda: tf.constant(1, dtype=tf.int32),
+    lambda: tf.constant(1, dtype=tf.float32))  # Error -- inconsistent dtypes: int32, float32
+```
+
+The same restriction in AutoGraph code:
+
+```
+if tf.random.uniform(()) > 0.5:
+  x = tf.constant(1, dtype=tf.int32)
+else:
+  x = tf.constant(1, dtype=tf.float32)  # Error -- inconsistent dtypes: int32, float32
+```
+
+Example of illegal dtype change in a loop:
+
+```
+# This won't work - "x" changes dtype inside the loop.
+x = tf.while_loop(
+    lambda _: tf.random.uniform(()) > 0.5,
+    lambda x: tf.constant(1, dtype=tf.float32),
+    loop_vars=(tf.constant(1, dtype=tf.int32),))  # Error -- inconsistent dtypes: int32, float32
+```
+
+The same restriction in AutoGraph code:
+
+```
+x = tf.constant(0, dtype=tf.int32)
+while tf.random.uniform(()) > 0.5:
+  x = tf.constant(0, dtype=tf.float32)   # Error -- inconsistent dtypes: int32, float32
+```
+
+#### Consistency of shape
+
+The shapes across all code paths must be consistent in loops only. When tensors
+do need to change shape across iterations, use `shape_invariants`.
+
+Note: Shapes are allowed to be inconsistent in conditionals. The result will be
+a partially dynamic shape.
+
+In a `tf.while_loop` (and correspondingly, an AutoGraph `while` or `for` loop)
+all loop variables must maintain consistent shape and dtype across iterations.
+That is, every loop variable must have the same shape at the end of the loop
+body as it had at the beginning of the loop body.
+
+Example of illegal shape change in a loop:
+
+```
+def loop_body(x):  # x.shape is ()
+  return tf.constant((1, 2, 3))  # Error -- inconsistent shapes: (), (3,)
+
+x = tf.while_loop(
+    lambda _: tf.random.uniform(()) > 0.5,
+    loop_body,
+    loop_vars=(tf.constant(1,))
+```
+
+The same restriction in AutoGraph code:
+
+```
+x = tf.constant(1,)
+while tf.random.uniform(()) > 0.5:
+  x = tf.constant((1, 2, 3))  # Error -- inconsistent shapes: (), (3,)
+```
+
+### Consistency of control flow types
+
+In AutoGraph, one can write Python control flow like `for i in range(10)`, as
+well as TensorFlow control flow like `for i in tf.range(10)`.
+
+However, one could also write (illegal) programs which start as Python control
+flow, then turn into TensorFlow control flow. In such cases, an error will be
+raised.
+
+Below are a few examples, along with recommendations.
+
+#### Python loop, TF-dependent break or return
+
+Example:
+
+```
+for i in range(10):
+  if tf.greater(i, 3):
+    break  # error - TF break inside Python loop
+```
+
+The solution in this case is to change the loop type to a TF loop:
+
+```
+for i in tf.range(10):
+  if tf.greater(i, 3):
+    break  # works
+```
+
+#### Python loop that turns into a TensorFlow loop
+
+Example:
+
+```
+i = 10
+while i > 0:
+  i = tf.math.subtract(i, 1)  # error - loop would turn into a TF loop
+```
+
+The solution in this case is to make sure the loop type starts as a TF loop,
+typically by making sure the condition is always a Tensor:
+
+```
+i = tf.constant(10)  # works
+while i > 0:
+  i = tf.math.subtract(i, 1)
+```
+
+#### TensorFlow loops never turn into Python loops
+
+Note that this is a legal case, as TensorFlow implicitly converts all Python
+values to Tensor:
+
+```
+i = tf.constant(10)
+while i > 0:
+  i = 0  # this is ok, will be auto-converted to Tensor
+```
+
+### Access to source code
+
+Key point: AutoGraph can only handle functions whose source code can be
+accessed at runtime.
+
+Almost all Python functions allow access to their source code. However, a few
+exceptions exist:
+
+ * functions created in the Python interactive shell
+ * functions with native bindings (these do not have Python source code)
+ * functions created dynamically, using `exec` or `eval`
+
+Use
+[inspect.findsource](https://docs.python.org/3/library/inspect.html#inspect.findsource)
+to quickly diagnose whether the source code is available for a function.
+
+For example:
+
+```
+import inspect
+
+def simple_function():
+  return 1
+
+# If this raises an error, then AutoGraph prints a warning.
+# If it returns source code, then AutoGraph should work as well.
+inspect.findsource(simple_function)
+```
+
+#### Source code of lambda functions
+
+##### TF 2.4 and newer
+
+Key Point: When nesting lambda functions, use distinguishing argument names
+to avoid parse errors.
+
+The Python runtime exposes the source code of lambda functions, however it
+may omit parts of the actual body, or include surrounding code. This may make it
+impossible to parse the exact source code of the lambda function (see
+https://github.com/tensorflow/tensorflow/issues/39832).
+
+AutoGraph uses alternate methods to parse the source code more robustly, but
+in rare cases it may be unable to distinguish between nested lambda functions
+of identical signatures.
+
+Example:
+
+```
+l = lambda x: lambda x: x + 1
+```
+
+AutoGraph raises an error for the code above because the parser cannot
+distinguish between the two function signatures. To work around this limitation,
+use distinct argument names:
+
+```
+l = lambda outer_x: lambda inner_x: inner_x + 1
+```
+
+##### Before TF 2.3 and older
+
+In older versions of TensorFlow, the loading code for lambda functions is not
+robust. Follow the guidance below to avoid errors.
+
+Important: Declare lambda functions on single lines to make sure their source
+code loads correctly.
+
+The Python runtime exposes the source code of lambda functions, however it
+may omit parts of the actual body, or include surrounding code. This may make it
+impossible to parse the exact source code of the lambda function.
+
+For example, consider the declaration of a lambda function below:
+
+```
+foo = (
+    lambda y: lambda x: x * y
+    - y
+)
+```
+
+The Python runtime will report the following source code for `foo`:
+
+```
+>>> inspect.getsource(foo)
+'    lambda y: lambda x: x*y \n'
+```
+
+In other cases, the source code it returns is not valid Python code, resulting
+in an error:
+
+```
+foo = (
+ 'bar',
+ lambda: x)
+```
+
+The reported source code contains an invalid token `)`:
+
+```
+>>> inspect.getsource(foo[1])
+' lambda: x)\n'
+```
+
+This shortcoming can be avoided by declaring the lambda in a single assignment
+or return value, and avoiding placing it inside parentheses which could cause
+auto-formatting tools to break it into multiple lines:
+
+
+```
+# Good - single assignment
+my_lambda = lambda: x
+
+# Good - single return
+return lambda x, y: x*y - y
+```
+
+```
+# Bad - wrapped in parentheses
+my_lambda = (lambda x, y: x * y - y)
+
+# Bad - inlined in another expression
+foo(lambda x, y: x + y, bar)
+```

tensorflow/python/autograph/g3doc/reference/operators.md

@@ -0,0 +1,422 @@
+# AutoGraph reference
+
+[Index](index.md)
+
+## Operator semantics
+
+### Definition
+
+This section describes the semantics of the operators used in code generated by
+AutoGraph. Understanding these operators will make it easier to read the
+generated code.
+
+AutoGraph operators are Python functions that replace certain Python constructs
+in the generated code.
+
+For example, the following statement:
+
+```
+if x:
+  y = 1
+else:
+  y = 2
+```
+
+Will result in the following generated code:
+
+```
+def get_state():
+    return (y,)
+
+def set_state(vars_):
+    nonlocal y
+    (y,) = vars_
+
+def if_body():
+    nonlocal y
+    y = 1
+
+def else_body():
+    nonlocal y
+    y = 2
+y = ag__.Undefined('y')
+ag__.if_stmt(ag__.ld(x), if_body, else_body, get_state, set_state, ('y',), 1)
+```
+
+In the example above, `ag__.if_stmt`, `ag__.ld` and `ag__.Undefined` are all
+AutoGraph operators.
+
+The source of truth for these operators is the [source code](https://github.com/tensorflow/tensorflow/tree/master/tensorflow/python/autograph/operators)
+. All public symbols exported by that module is considered an operator.
+
+### Type-based dispatch
+
+AutoGraph replaces Python statements with operators in order to enable
+type-based dispatch. If Python didn't support things like `__add__`, then
+AutoGraph would already have an `add` operator.
+
+Dispatch means simply that the operator does different things based on the type
+of input.
+
+Generally, the dispatch follows these rules:
+ * if the input is a type that would execute normally under Python (this is also
+   referred to as "the default path"), then AutoGraph always reverts to the
+   corresponding Python operator. For example, `ag__.not(False)` always has the
+   same result as `not False`.
+ * if the input is a TensorFlow type, then AutoGraph typically dispatches to an
+   equivalent TensorFlow API, performs additional checks or just raises an
+   error. For example, `ag__.eq(tf.constant(1), tf.constant(2))` has the same
+   result as `tf.math.equal(tf.constant(1), tf.constant(2))`.
+
+The first rule above means that if you convert normal, non-TensorFlow code with
+AutoGraph and call it with non-TensorFlow inputs, executing the generated code
+should be no different than executing the original.
+
+### Functional form
+
+All AutoGraph operators use pure functional forms. This may sometimes mean that
+expressions which normally appear bare in Python, are wrapped inside a function
+(also known as thunk). If a Python statement appears as just `foo`, then a
+corresponding thunk is `lambda: foo`.
+
+### Operator list
+
+#### Conditional expressions
+
+[Source](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/autograph/operators/conditional_expressions.py)
+
+##### `if_exp`
+
+[Source](https://github.com/tensorflow/tensorflow/blob/bacd16a95d5a6f3d5081e3d56c515671c784d289/tensorflow/python/autograph/operators/conditional_expressions.py#L27)
+
+The Python conditional statement: `foo if bar else baz`.
+
+Args: cond: expression condition; same as `cond` in `_ if cond else _`. if_true:
+true value (as thunk); same as `lambda: x` in `x if _ else _`. if_false: false
+value (as thunk); same as `lambda: x` in `_ if _ else x`. expr_repr:
+human-readable string representing `cond`. Used for error messages.
+
+Example:
+
+```
+true_val if cond else false_val
+```
+
+```
+ag__.if_expr(cond, lambda: true_val, lambda: false_val, 'cond')
+```
+
+Dispatch on `cond`:
+
+*   default: to Python if-else statement.
+*   tf.Tensor: to `tf.cond`, checking that `true_val` and `false_val` have
+    compatible shape and type.
+
+#### Control flow
+
+[Source](https://github.com/tensorflow/tensorflow/blob/bacd16a95d5a6f3d5081e3d56c515671c784d289/tensorflow/python/autograph/operators/control_flow.py)
+
+Unlike Python, AutoGraph control flow operators use explicit control flow
+variables, which include all symbols which are modified by the control flow.
+
+For example, the code below has a single loop variable, `x`:
+
+```
+while x < 3:
+  x = x + 1
+```
+
+In addition, control flow that is dispatched to non-Python implementation is
+subject to restrictions of the respective implementations. For example,
+`tf.while_loop` requires that all loop variables have supported types (e.g.
+`Tensor` of consistent shape and dtype).
+
+##### `for_stmt`
+
+[Source](https://github.com/tensorflow/tensorflow/blob/bacd16a95d5a6f3d5081e3d56c515671c784d289/tensorflow/python/autograph/operators/control_flow.py#L369)
+
+For loop: `for var in target: body`, extended with a per-iteration
+condition to handle early termination (e.g. due to a `break`).
+
+Args:
+
+*   iter_: iteration target; same as `n` in `for _ in n`.
+*   extra_test: optional extra per-iteration condition (as thunk).
+*   body: loop body (as unary thunk); same as `def body(i): <b>` in `for i in _:
+    <b>`.
+*   get_state: returns the current value of the loop variables
+*   set_state: sets new values into the loop variables
+*   symbol_names: human-readable string representing each loop variable. Used
+    for error messages.
+*   opts: additional, implementation-specific, keyword arguments.
+
+Example:
+
+```
+for i in range(3):
+  j = j + i
+```
+
+```
+def get_state():
+    return (j,)
+
+def set_state(vars_):
+    nonlocal j
+    (j,) = vars_
+
+def loop_body(itr):
+    nonlocal j
+    i = itr
+    j = j + i
+
+ag__.for_stmt(range(3), None, loop_body, get_state, set_state, ('j',), {})
+```
+
+Example (using extra_test):
+
+```
+for i in range(3):
+  if i > 2:
+    break
+  j = j + i
+```
+
+```
+def get_state():
+    return (j,)
+
+def set_state(vars_):
+    nonlocal j
+    (j,) = vars_
+
+def loop_body(itr):
+    nonlocal j
+    i = itr
+    j = j + i
+
+def extra_test():
+    return not(i <= 2)
+
+ag__.for_stmt(range(3), extra_test, loop_body, get_state, set_state, ('j',), {})
+```
+
+Dispatch on `iter_`:
+
+*   default: to Python for loop (accounting for `extra_test`).
+*   `tf.Tensor` produced by `tf.range`: to `tf.while_loop`, removing the
+    `tf.range`.
+*   `tf.Tensor`, `tf.RagedTensor`: to `tf.while_loop`, checking the loop vars
+    for consistency. `opts` forwarded to `tf.while_loop`. Iterates over the
+    outermost dimension of the tensor (similar to `tf.map_fn`).
+*   `tf.data.Dataset`: to `tf.data.Dataset.take_while`, checking the loop vars
+    for consistency.
+*   `tf.data.Iterator`, `tf.distribute.Iterator`: to `tf.while_loop` called on
+    the iterator's `get_next_as_optional`, checking the loop vars for
+    consistency.
+*   `tf.distribute.Iterable`: to `tf.distribute.Iterable.reduce`.
+
+##### `if_stmt`
+
+[Source](https://github.com/tensorflow/tensorflow/blob/bacd16a95d5a6f3d5081e3d56c515671c784d289/tensorflow/python/autograph/operators/control_flow.py#L1125)
+
+If statement: `if cond: body else: orelse-body`.
+
+Args:
+
+*   cond: if condition; same as `cond` in `if cond`.
+*   body: true branch (as unary thunk); same as `def body(): <b>` in `if _:
+    <b>`.
+*   orelse: false branch (as unary thunk); same as `def body(): <b>` in `if _:
+    <b>`.
+*   get_state: returns the current value of the conditional variables
+*   set_state: sets new values into the conditional variables
+*   symbol_names: human-readable string representing each conditional variable.
+    Used for error messages.
+*   nouts: number of output conditional variables. Not all conditional variables
+    are outputs - some are just inputs. The first nouts values in get_state and
+    set_state are the conditional outputs.
+
+Example:
+
+```
+if k > 1:
+  j = j + i
+```
+
+```
+def get_state():
+    return (j, i)
+
+def set_state(vars_):
+    nonlocal j, i
+    (j, i) = vars_
+
+def body():
+    nonlocal j, i
+    j = j + i
+
+def orelse():
+    pass
+
+ag__.if_stmt(k > 1, body, orelse, get_state, set_state, ('j', 'i'), 1)
+```
+
+Dispatch on `cond`:
+
+*   default: to Python if statement.
+*   `tf.Tensor`: to `tf.cond`, removing the `tf.range`.
+
+##### `while_stmt`
+
+[Source](https://github.com/tensorflow/tensorflow/blob/bacd16a95d5a6f3d5081e3d56c515671c784d289/tensorflow/python/autograph/operators/control_flow.py#L811)
+
+While loop: `while cond: body`.
+
+Args:
+
+*   test: loop condition (as thunk); same as `def test(): cond` in `while cond`.
+*   body: loop body (as thunk); same as `def body(): <b>` in `while _: <b>`.
+*   get_state: returns the current value of the loop variables
+*   set_state: sets new values into the loop variables
+*   symbol_names: human-readable string representing each loop variable. Used
+    for error messages.
+*   opts: additional, implementation-specific, keyword arguments.
+
+Example:
+
+```
+while j > 10:
+  j = j + i
+```
+
+```
+def get_state():
+    return (j,)
+
+def set_state(vars_):
+    nonlocal j
+    (j,) = vars_
+
+def loop_test():
+    nonlocal j
+    return j > 10
+
+def loop_body():
+    nonlocal j
+    j = j + i
+
+ag__.while_stmt(loop_test, loop_body, get_state, set_state, ('j',), {})
+```
+
+Dispatch on return type of `test`:
+
+*   default: to Python while loop.
+*   `tf.Tensor`: to `tf.while_loop`.
+
+#### Data structures
+
+[Source](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/autograph/operators/data_structures.py)
+
+##### `list_append`
+
+[Source](https://github.com/tensorflow/tensorflow/blob/bacd16a95d5a6f3d5081e3d56c515671c784d289/tensorflow/python/autograph/operators/data_structures.py#L171)
+
+List append operation: `l.append(x)`. Callers should assume that the list
+argument is modified, if that is possible.
+
+Args:
+
+*   list_: a list-like value.
+*   x: value to append to list.
+
+Returns:
+
+*   same as list_, with an appended value.
+
+Example:
+
+```
+l.append(x)
+```
+
+```
+l = ag__.list_append(l, x)
+```
+
+Dispatch on `list_`:
+
+*   default: to `list_.append`.
+*   `tf.Tensor`: to `tf.raw_ops.tensor_list_push_back`.
+*   `tf.TensorArray`: to `tf.TensorArray.write`.
+
+##### `list_pop`
+
+[Source](https://github.com/tensorflow/tensorflow/blob/bacd16a95d5a6f3d5081e3d56c515671c784d289/tensorflow/python/autograph/operators/data_structures.py#L235)
+
+List pop operation: `l.pop(i)`. Callers should assume that the list
+argument is modified, if that is possible.
+
+Args:
+
+*   list_: a list-like value.
+*   i: optional index to remove from.
+*   opts: optional, implementation-specific arguments.
+
+Returns:
+
+*   new_list: same as list_, with the value removed
+*   x: the value that was removed
+
+Example:
+
+```
+x = l.pop()
+```
+
+```
+l, x = ag__.list_pop(l)
+```
+
+Dispatch on `list_`:
+
+*   default: to `list_.pop`.
+*   `tf.Tensor`: to `tf.raw_ops.tensor_list_pop_back`.
+
+##### `list_stack`
+##### `ListPopOpts`
+##### `ListStackOpts`
+##### `new_list`
+
+### Exceptions
+
+##### `assert_stmt`
+
+### Boolean
+
+##### `and_`
+##### `eq`
+##### `not_`
+##### `not_eq`
+##### `or_`
+
+### Python built-ins
+
+##### `float_`
+##### `int_`
+##### `len_`
+##### `print_`
+##### `range_`
+
+### Slicing
+
+##### `get_item`
+##### `GetItemOpts`
+##### `set_item`
+
+### Variables
+
+##### `ld`
+##### `ldu`
+##### `Undefined`
+##### `UndefinedReturnValue`

tensorflow/python/autograph/impl/BUILD

@@ -0,0 +1,116 @@
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:tensorflow.default.bzl", "tf_py_strict_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "api",
+    srcs = ["api.py"],
+    strict_deps = True,
+    visibility = ["//tensorflow:__subpackages__"],
+    deps = [
+        ":conversion",
+        "//tensorflow/python/autograph/converters:asserts",
+        "//tensorflow/python/autograph/converters:break_statements",
+        "//tensorflow/python/autograph/converters:call_trees",
+        "//tensorflow/python/autograph/converters:conditional_expressions",
+        "//tensorflow/python/autograph/converters:continue_statements",
+        "//tensorflow/python/autograph/converters:control_flow",
+        "//tensorflow/python/autograph/converters:directives",
+        "//tensorflow/python/autograph/converters:functions",
+        "//tensorflow/python/autograph/converters:lists",
+        "//tensorflow/python/autograph/converters:logical_expressions",
+        "//tensorflow/python/autograph/converters:return_statements",
+        "//tensorflow/python/autograph/converters:slices",
+        "//tensorflow/python/autograph/converters:variables",
+        "//tensorflow/python/autograph/core:ag_ctx",
+        "//tensorflow/python/autograph/core:converter",
+        "//tensorflow/python/autograph/core:function_wrappers",
+        "//tensorflow/python/autograph/core:unsupported_features_checker",
+        "//tensorflow/python/autograph/lang:special_functions",
+        "//tensorflow/python/autograph/operators:__init__",
+        "//tensorflow/python/autograph/operators:py_builtins",
+        "//tensorflow/python/autograph/pyct:anno",
+        "//tensorflow/python/autograph/pyct:cfg",
+        "//tensorflow/python/autograph/pyct:error_utils",
+        "//tensorflow/python/autograph/pyct:errors",
+        "//tensorflow/python/autograph/pyct:inspect_utils",
+        "//tensorflow/python/autograph/pyct:origin_info",
+        "//tensorflow/python/autograph/pyct:qual_names",
+        "//tensorflow/python/autograph/pyct:transpiler",
+        "//tensorflow/python/autograph/pyct/static_analysis:activity",
+        "//tensorflow/python/autograph/pyct/static_analysis:reaching_definitions",
+        "//tensorflow/python/autograph/utils:__init__",
+        "//tensorflow/python/autograph/utils:ag_logging",
+        "//tensorflow/python/eager/polymorphic_function:tf_method_target",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/util:tf_decorator_py",
+        "//tensorflow/python/util:tf_export",
+        "//tensorflow/python/util:tf_inspect",
+        "//tensorflow/python/util:tf_stack",
+    ],
+)
+
+py_library(
+    name = "conversion",
+    srcs = ["conversion.py"],
+    strict_deps = True,
+    visibility = ["//visibility:private"],
+    deps = [
+        "//tensorflow/python/autograph/core:config",
+        "//tensorflow/python/autograph/pyct:cache",
+        "//tensorflow/python/autograph/pyct:inspect_utils",
+        "//tensorflow/python/autograph/utils:ag_logging",
+        "//tensorflow/python/eager/polymorphic_function:tf_method_target",
+        "//tensorflow/python/util:tf_inspect",
+    ],
+)
+
+tf_py_strict_test(
+    name = "api_test",
+    srcs = ["api_test.py"],
+    deps = [
+        ":api",
+        ":conversion",
+        "//tensorflow/python/autograph/core:ag_ctx",
+        "//tensorflow/python/autograph/core:converter",
+        "//tensorflow/python/autograph/core:test_lib",
+        "//tensorflow/python/autograph/pyct:errors",
+        "//tensorflow/python/autograph/pyct:inspect_utils",
+        "//tensorflow/python/autograph/pyct:parser",
+        "//tensorflow/python/autograph/utils:ag_logging",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:_errors_test_helper",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/util:function_utils",
+        "//tensorflow/python/util:tf_decorator_py",
+        "//tensorflow/python/util:tf_inspect",
+        "//third_party/py/numpy",
+    ],
+)
+
+tf_py_strict_test(
+    name = "conversion_test",
+    srcs = ["conversion_test.py"],
+    deps = [
+        ":api",
+        ":conversion",
+        "//tensorflow/python/autograph/core:config",
+        "//tensorflow/python/autograph/core:converter",
+        "//tensorflow/python/autograph/impl/testing:pybind_for_testing",
+        "//tensorflow/python/autograph/utils:__init__",
+        "//tensorflow/python/eager:function",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)

tensorflow/python/autograph/impl/api.py

@@ -0,0 +1,949 @@
+# Copyright 2016 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.
+# ==============================================================================
+"""This module contains the user- and codegen-facing API for AutoGraph."""
+
+import functools
+import importlib
+import inspect
+import os
+import sys
+import textwrap
+import traceback
+
+from tensorflow.python.autograph import operators
+from tensorflow.python.autograph import utils
+from tensorflow.python.autograph.converters import asserts
+from tensorflow.python.autograph.converters import break_statements
+from tensorflow.python.autograph.converters import call_trees
+from tensorflow.python.autograph.converters import conditional_expressions
+from tensorflow.python.autograph.converters import continue_statements
+from tensorflow.python.autograph.converters import control_flow
+from tensorflow.python.autograph.converters import directives
+from tensorflow.python.autograph.converters import functions
+from tensorflow.python.autograph.converters import lists
+from tensorflow.python.autograph.converters import logical_expressions
+from tensorflow.python.autograph.converters import return_statements
+from tensorflow.python.autograph.converters import slices
+from tensorflow.python.autograph.converters import variables
+from tensorflow.python.autograph.core import ag_ctx
+from tensorflow.python.autograph.core import converter
+from tensorflow.python.autograph.core import function_wrappers
+from tensorflow.python.autograph.core import unsupported_features_checker
+from tensorflow.python.autograph.impl import conversion
+from tensorflow.python.autograph.lang import special_functions
+from tensorflow.python.autograph.operators import py_builtins
+from tensorflow.python.autograph.pyct import anno
+from tensorflow.python.autograph.pyct import cfg
+from tensorflow.python.autograph.pyct import error_utils
+from tensorflow.python.autograph.pyct import errors
+from tensorflow.python.autograph.pyct import inspect_utils
+from tensorflow.python.autograph.pyct import origin_info
+from tensorflow.python.autograph.pyct import qual_names
+from tensorflow.python.autograph.pyct import transpiler
+from tensorflow.python.autograph.pyct.static_analysis import activity
+from tensorflow.python.autograph.pyct.static_analysis import reaching_definitions
+from tensorflow.python.autograph.utils import ag_logging as logging
+from tensorflow.python.eager.polymorphic_function import tf_method_target
+from tensorflow.python.framework import errors_impl
+from tensorflow.python.util import tf_decorator
+from tensorflow.python.util import tf_inspect
+from tensorflow.python.util import tf_stack
+from tensorflow.python.util.tf_export import tf_export
+
+
+def is_autograph_strict_conversion_mode():
+  return int(os.environ.get('AUTOGRAPH_STRICT_CONVERSION', '0')) > 0
+
+
+#
+# Error handling
+#
+
+
+# TODO(mdan): Export this symbol.
+class AutoGraphError(errors.PyCTError):
+  """Base class for all AutoGraph exceptions."""
+  pass
+
+
+class ConversionError(AutoGraphError):
+  """Raised during the conversion process."""
+  pass
+
+
+class StagingError(AutoGraphError):
+  """Raised during the staging (i.e. Python execution) of converted code."""
+  pass
+
+
+class _ErrorMetadata(error_utils.ErrorMetadataBase):
+  """AutoGraph-specific error metadata. See base class."""
+
+  def create_exception(self, source_error):
+    preferred_type = type(source_error)
+    if issubclass(preferred_type, errors_impl.OpError):
+      # Best-effort unpacking of OpError exceptions.
+      # TODO(mdan): Use a mechanism that is more future-proof.
+      init_argspec = tf_inspect.getfullargspec(preferred_type.__init__)
+      message = self.get_message()
+      init_args = tuple(init_argspec.args)
+      # At the time of this writing, TF errors either take 3 or 4 arguments,
+      # the argument '*args' may or may not be used.
+      if init_args == ('self', 'node_def', 'op', 'message'):
+        return preferred_type(source_error.node_def, source_error.op, message,
+                              source_error.experimental_payloads)
+
+    elif preferred_type in (errors.PyCTError, AutoGraphError, ConversionError,
+                            StagingError, errors_impl.InaccessibleTensorError,
+                            errors_impl.OperatorNotAllowedInGraphError):
+      return preferred_type(self.get_message())
+
+    exc = super(_ErrorMetadata, self).create_exception(source_error)
+    if exc is not None:
+      return exc
+
+    # Note: While changing an error's message property to change the message it
+    # displays will probably work a lot of times, there is no standard way in
+    # Python to do that. The safest way is therefore to create a new exception.
+    # For user defined exceptions, we could define an interface that allowed
+    # them to work under this mechanism.
+    return StagingError(self.get_message())
+
+
+def _attach_error_metadata(e, f):
+  """Augments an error with the metadata necessary for rewrite."""
+  if hasattr(e, 'ag_pass_through'):
+    return
+
+  metadata = getattr(e, 'ag_error_metadata', None)
+  source_map = f.ag_source_map
+
+  if metadata is None:
+    logging.log(1, 'Caught error in user callable %s', f, exc_info=True)
+    message = '{}: {}'.format(e.__class__.__name__, e)
+  else:
+    message = None
+
+  cause_tb = traceback.extract_tb(sys.exc_info()[2])[1:]
+
+  e.ag_error_metadata = _ErrorMetadata(cause_tb, metadata, message, source_map,
+                                       __file__)
+
+
+class StackTraceMapper(tf_stack.StackTraceMapper):
+  """Remaps generated code to code it originated from."""
+
+  def __init__(self, converted_fn):
+    super().__init__()
+    self._source_map = converted_fn.ag_source_map
+    # This may be called repeatedly: once on entry, by the superclass, then by
+    # each child context manager.
+    self._cached_map = None
+
+  def get_effective_source_map(self):
+    if self._cached_map is not None:
+      return self._cached_map
+
+    parent_map = self.parent.get_effective_source_map()
+
+    effective_source_map = {}
+    for loc, origin in self._source_map.items():
+      effective_source_map[(loc.filename, loc.lineno)] = (origin.loc.filename,
+                                                          origin.loc.lineno,
+                                                          origin.function_name)
+
+    for key, value in parent_map.items():
+      filename, lineno, _ = value
+      value_loc = origin_info.LineLocation(filename=filename, lineno=lineno)
+      if value_loc in self._source_map:
+        origin = self._source_map[value_loc]
+        effective_source_map[key] = (origin.loc.filename, origin.loc.lineno,
+                                     origin.function_name)
+      else:
+        effective_source_map[key] = value
+
+    self._cached_map = effective_source_map
+    return effective_source_map
+
+
+#
+# Actual source code transformation
+#
+
+
+class PyToTF(transpiler.PyToPy):
+  """The TensorFlow AutoGraph transformer."""
+
+  def __init__(self):
+    super(PyToTF, self).__init__()
+    self._extra_locals = None
+
+  def get_transformed_name(self, node):
+    return 'tf__' + super(PyToTF, self).get_transformed_name(node)
+
+  def get_extra_locals(self):
+    if self._extra_locals is None:
+      # TODO(mdan): Move into core or replace with an actual importable module.
+      # Craft a module that exposes the external API as well as certain
+      # internal modules.
+      module_spec = importlib.machinery.ModuleSpec('autograph', None)
+      ag_internal = importlib.util.module_from_spec(module_spec)
+      ag_internal.__dict__.update(inspect.getmodule(PyToTF).__dict__)
+      ag_internal.ConversionOptions = converter.ConversionOptions
+      ag_internal.STD = converter.STANDARD_OPTIONS
+      ag_internal.Feature = converter.Feature
+      ag_internal.utils = utils
+      ag_internal.FunctionScope = function_wrappers.FunctionScope
+      ag_internal.with_function_scope = function_wrappers.with_function_scope
+      # TODO(mdan): Add safeguards against name clashes.
+      # We don't want to create a submodule because we want the operators to be
+      # accessible as ag__.<operator>
+      ag_internal.__dict__.update(special_functions.__dict__)
+      ag_internal.__dict__.update(operators.__dict__)
+
+      self._extra_locals = {'ag__': ag_internal}
+    return self._extra_locals
+
+  def get_caching_key(self, ctx):
+    return ctx.options
+
+  def initial_analysis(self, node, ctx):
+    graphs = cfg.build(node)
+    node = qual_names.resolve(node)
+    node = activity.resolve(node, ctx, None)
+    node = reaching_definitions.resolve(node, ctx, graphs)
+    anno.dup(
+        node,
+        {
+            anno.Static.DEFINITIONS: anno.Static.ORIG_DEFINITIONS,
+        },
+    )
+    return node
+
+  def transform_ast(self, node, ctx):
+    unsupported_features_checker.verify(node)
+    node = self.initial_analysis(node, ctx)
+
+    node = functions.transform(node, ctx)
+    node = directives.transform(node, ctx)
+    node = break_statements.transform(node, ctx)
+    if ctx.user.options.uses(converter.Feature.ASSERT_STATEMENTS):
+      node = asserts.transform(node, ctx)
+    # Note: sequencing continue canonicalization before for loop one avoids
+    # dealing with the extra loop increment operation that the for
+    # canonicalization creates.
+    node = continue_statements.transform(node, ctx)
+    node = return_statements.transform(node, ctx)
+    if ctx.user.options.uses(converter.Feature.LISTS):
+      node = lists.transform(node, ctx)
+      node = slices.transform(node, ctx)
+    node = call_trees.transform(node, ctx)
+    node = control_flow.transform(node, ctx)
+    node = conditional_expressions.transform(node, ctx)
+    node = logical_expressions.transform(node, ctx)
+    node = variables.transform(node, ctx)
+    return node
+
+
+def _convert_actual(entity, program_ctx):
+  """Applies AutoGraph to entity."""
+
+  # TODO(mdan): Put these extra fields inside __autograph_info__.
+  if not hasattr(entity, '__code__'):
+    raise ValueError('Cannot apply autograph to a function that doesn\'t '
+                     'expose a __code__ object. If this is a @tf.function,'
+                     ' try passing f.python_function instead.')
+
+  transformed, module, source_map = _TRANSPILER.transform(entity, program_ctx)
+
+  assert not hasattr(transformed, 'ag_module')
+  assert not hasattr(transformed, 'ag_source_map')
+  transformed.ag_module = module
+  transformed.ag_source_map = source_map
+  return transformed
+
+
+#
+# Generated code support
+#
+
+
+def autograph_artifact(entity, extras=None):
+  if inspect.ismethod(entity):
+    setattr(entity.__func__, 'autograph_info__', extras)
+  else:
+    setattr(entity, 'autograph_info__', extras)
+  return entity
+
+
+def is_autograph_artifact(entity):
+  return hasattr(entity, 'autograph_info__')
+
+
+def converted_call(f, args, kwargs, caller_fn_scope=None, options=None):
+  """Converts a function call inline.
+
+  For internal use only.
+
+  Note: The argument list is optimized for readability of generated code, which
+  may look like this:
+
+    ag__.converted_call(f, (arg1, arg2), None, fscope)
+    ag__.converted_call(f, (), dict(arg1=val1, **kwargs), fscope)
+    ag__.converted_call(f, (arg1, arg2) + varargs, dict(**kwargs), lscope)
+
+  Args:
+    f: The function to convert.
+    args: Tuple, the original positional arguments of f
+    kwargs: Optional[Dict], the original keyword arguments of f
+    caller_fn_scope: Optional[function_wrappers.FunctionScope], the function
+      scope of the converted function in which this call was originally made.
+    options: Optional[converter.ConversionOptions], conversion options. If not
+      specified, the value of caller_fn_scope.callopts is used. Either options
+      or caller_fn_scope must be present.
+
+  Returns:
+    Any, the result of executing a possibly-converted `f` with the given
+      arguments.
+  """
+  logging.log(1, 'Converted call: %s\n    args: %s\n    kwargs: %s\n', f, args,
+              kwargs)
+
+  if options is None:
+    if caller_fn_scope is None:
+      raise ValueError('either caller_fn_scope or options must have a value')
+    options = caller_fn_scope.callopts
+
+  if conversion.is_in_allowlist_cache(f, options):
+    logging.log(2, 'Allowlisted %s: from cache', f)
+    return _call_unconverted(f, args, kwargs, options, False)
+
+  if ag_ctx.control_status_ctx().status == ag_ctx.Status.DISABLED:
+    logging.log(2, 'Allowlisted: %s: AutoGraph is disabled in context', f)
+    return _call_unconverted(f, args, kwargs, options, False)
+
+  if is_autograph_artifact(f):
+    logging.log(2, 'Permanently allowed: %s: AutoGraph artifact', f)
+    return _call_unconverted(f, args, kwargs, options)
+
+  # If this is a partial, unwrap it and redo all the checks.
+  if isinstance(f, functools.partial):
+    new_kwargs = {}
+    if f.keywords is not None:
+      # Use copy to avoid mutating the underlying keywords.
+      new_kwargs = f.keywords.copy()
+    if kwargs is not None:
+      new_kwargs.update(kwargs)
+    new_args = f.args + args
+    logging.log(3, 'Forwarding call of partial %s with\n%s\n%s\n', f, new_args,
+                new_kwargs)
+    return converted_call(
+        f.func,
+        new_args,
+        new_kwargs,
+        caller_fn_scope=caller_fn_scope,
+        options=options)
+
+  if inspect_utils.isbuiltin(f):
+    if f is eval:
+      return py_builtins.eval_in_original_context(f, args, caller_fn_scope)
+    if f is super:
+      return py_builtins.super_in_original_context(f, args, caller_fn_scope)
+    if f is globals:
+      return py_builtins.globals_in_original_context(caller_fn_scope)
+    if f is locals:
+      return py_builtins.locals_in_original_context(caller_fn_scope)
+    if kwargs:
+      return py_builtins.overload_of(f)(*args, **kwargs)
+    else:
+      return py_builtins.overload_of(f)(*args)
+
+  if conversion.is_unsupported(f):
+    return _call_unconverted(f, args, kwargs, options)
+
+  if not options.user_requested and conversion.is_allowlisted(f):
+    return _call_unconverted(f, args, kwargs, options)
+
+  # internal_convert_user_code is for example turned off when issuing a dynamic
+  # call conversion from generated code while in nonrecursive mode. In that
+  # case we evidently don't want to recurse, but we still have to convert
+  # things like builtins.
+  if not options.internal_convert_user_code:
+    return _call_unconverted(f, args, kwargs, options)
+
+  try:
+    if inspect.ismethod(f) or inspect.isfunction(f):
+      target_entity = f
+      effective_args = args
+
+      f_self = getattr(f, '__self__', None)
+      if f_self is not None:
+        if isinstance(f_self, tf_method_target.TfMethodTarget):
+          f_self = f_self.target
+        effective_args = (f_self,) + effective_args
+
+    elif hasattr(f, '__class__') and hasattr(f.__class__, '__call__'):
+      # Callable objects. Dunder methods have special lookup rules, see:
+      # https://docs.python.org/3/reference/datamodel.html#specialnames
+      # TODO(mdan): Recurse into converted_call to simplify other verifications.
+      # This should be handled in the same way as partials.
+      target_entity = f.__class__.__call__
+      effective_args = (f,) + args
+
+    else:
+      target_entity = f
+      raise NotImplementedError('unknown callable type "%s"' % type(f))
+
+  except Exception as e:  # pylint:disable=broad-except
+    logging.log(1, 'Error transforming entity %s', target_entity, exc_info=True)
+    if is_autograph_strict_conversion_mode():
+      raise
+    return _fall_back_unconverted(f, args, kwargs, options, e)
+
+  if not hasattr(target_entity, '__code__'):
+    logging.log(2, 'Permanently allowed: %s: native binding', target_entity)
+    return _call_unconverted(f, args, kwargs, options)
+  elif (hasattr(target_entity.__code__, 'co_filename') and
+        target_entity.__code__.co_filename == '<string>'):
+    # TODO(mdan): __globals__['txt'] might work in Py3.
+    logging.log(2, 'Permanently allowed: %s: dynamic code (exec?)',
+                target_entity)
+    return _call_unconverted(f, args, kwargs, options)
+
+  try:
+    program_ctx = converter.ProgramContext(options=options)
+    converted_f = _convert_actual(target_entity, program_ctx)
+    if logging.has_verbosity(2):
+      _log_callargs(converted_f, effective_args, kwargs)
+  except Exception as e:  # pylint:disable=broad-except
+    logging.log(1, 'Error transforming entity %s', target_entity, exc_info=True)
+    if is_autograph_strict_conversion_mode():
+      raise
+    return _fall_back_unconverted(f, args, kwargs, options, e)
+
+  with StackTraceMapper(converted_f), tf_stack.CurrentModuleFilter():
+    try:
+      if kwargs is not None:
+        result = converted_f(*effective_args, **kwargs)
+      else:
+        result = converted_f(*effective_args)
+    except Exception as e:
+      _attach_error_metadata(e, converted_f)
+      raise
+
+  return result
+
+
+def _call_unconverted(f, args, kwargs, options, update_cache=True):
+  """Calls the original function without converting with AutoGraph."""
+  if update_cache:
+    conversion.cache_allowlisted(f, options)
+
+  if (inspect.ismethod(f) and
+      isinstance(f.__self__, tf_method_target.TfMethodTarget)):
+    return f.__self__.call(args, kwargs)
+
+  if kwargs is not None:
+    return f(*args, **kwargs)
+  return f(*args)
+
+
+def _fall_back_unconverted(f, args, kwargs, options, exc):
+  """Falls back to calling the function unconverted, in case of error."""
+  # TODO(mdan): Consider adding an internal metric.
+  warning_template = (
+      'AutoGraph could not transform %s and will run it as-is.\n'
+      '%s'
+      'Cause: %s\n'
+      'To silence this warning, decorate the function with'
+      ' @tf.autograph.experimental.do_not_convert')
+  if isinstance(exc, errors.InaccessibleSourceCodeError):
+    if ag_ctx.INSPECT_SOURCE_SUPPORTED:
+      logging.warning(warning_template, f, '', exc)
+  elif isinstance(exc, errors.UnsupportedLanguageElementError):
+    if not conversion.is_in_allowlist_cache(f, options):
+      logging.warning(warning_template, f, '', exc)
+  else:
+    file_bug_message = (
+        'Please report this to the TensorFlow team. When filing the bug, set'
+        ' the verbosity to 10 (on Linux, `export AUTOGRAPH_VERBOSITY=10`) and'
+        ' attach the full output.\n')
+    logging.warning(warning_template, f, file_bug_message, exc)
+
+  return _call_unconverted(f, args, kwargs, options)
+
+
+#
+# TensorFlow integration
+#
+
+
+@tf_export('__internal__.autograph.tf_convert', v1=[])
+def tf_convert(f, ctx, convert_by_default=True, user_requested=False):
+  """Decorator that applies AutoGraph to a function.
+
+  Use in internal APIs.
+
+  This API is suitable for high order functions internal to the TensorFlow API,
+  and more generally any function to which AutoGraph is not applied.
+
+  Guidance: `convert` was a decorator meant for use directly by developers, but
+  most of today's uses go through `tf.function`. `tf_convert` is to be called
+  from high order functions internal to TF. By default, all the internal
+  TensorFlow functions are skipped when AutoGraph processes the code. This may
+  lead to user-supplied functions to be incorrectly skipped as well.
+  `tf_convert` helps avoid that. See the following example for more details.
+
+  ```
+  =====tf_internal_module.py=====
+
+  def unconverted(input_fn):
+    return input_fn()
+
+  def converted(input_fn):
+    return tf.__internal__.autograph.tf_convert(
+       input_fn, ctx=tf.__internal__.autograph.control_status_ctx())()
+
+  ======user_module.py======
+
+  @tf.function
+  def foo(input_fn)
+    return unconverted(input_fn)
+
+  @tf.function
+  def bar(input_fn)
+    return converted(input_fn)
+
+  @tf.function(autograph=False)
+  def baz(input_fn)
+    return converted(input_fn)
+  ```
+
+  The `foo` method above will execute the `input_fn` without autograph
+  conversion, while the `bar` method will run an autographed `input_fn`. The
+  `baz` method will run an unconverted `input_fn`, since `tf_convert` respect
+  the control status context.
+
+  Note that both methods in `tf_internal_module` are skipped by autograph when
+  tracing the `tf.function`. The configuration of whether a module/package
+  should be skipped by autograph is controlled in
+  tensorflow/python/autograph/core/config.py.
+
+  Args:
+    f: Callable.
+    ctx: ag_ctx.ControlStatusCtx, the Autograph context in which `f` is used.
+    convert_by_default: bool, whether to use AutoGraph when the context doesn't
+      specify.
+    user_requested: bool, whether to ignore the conversion allowlist. See
+      ConversionOptions.user_requested.
+
+  Returns:
+    Either `f or the converted version of `f`.
+  """
+
+  if is_autograph_artifact(f):
+    return f
+  f_wrapper = f
+  decorators, f = tf_decorator.unwrap(f)
+
+  # TODO(mdan): Grab features from context.
+  # Note: we pass the original context through to convert to properly handle the
+  # following scenario, which can be used inside TF implementations:
+  #
+  #   ctx = ag_ctx.control_status_ctx()
+  #   @function(autograph=False)  # Low-level graph code
+  #   def inner_fn():
+  #     # The context is disabled here, but should be enabled in user user_fn
+  #     tf_convert(user_fn, ctx=ctx)
+  if ctx.status == ag_ctx.Status.ENABLED:
+    wrapper_factory = convert(
+        recursive=True, user_requested=user_requested, conversion_ctx=ctx)
+  elif ctx.status == ag_ctx.Status.DISABLED:
+    wrapper_factory = do_not_convert
+  elif ctx.status == ag_ctx.Status.UNSPECIFIED:
+    if convert_by_default:
+      wrapper_factory = convert(
+          recursive=True, user_requested=user_requested, conversion_ctx=ctx)
+    else:
+      wrapper_factory = call_with_unspecified_conversion_status
+  else:
+    assert False, 'This switch contains all possible cases!'
+  wrapper = wrapper_factory(f)
+
+  if decorators:
+    wrapper = tf_decorator.rewrap(f_wrapper, f, wrapper)
+
+  return autograph_artifact(wrapper)
+
+
+def call_with_unspecified_conversion_status(func):
+  """Decorator that resets the conversion context to the unspecified status."""
+
+  def wrapper(*args, **kwargs):
+    with ag_ctx.ControlStatusCtx(status=ag_ctx.Status.UNSPECIFIED):
+      return func(*args, **kwargs)
+
+  if inspect.isfunction(func) or inspect.ismethod(func):
+    wrapper = functools.update_wrapper(wrapper, func)
+
+  return autograph_artifact(wrapper)
+
+
+def _log_callargs(f, args, kwargs):
+  """Logging helper."""
+  logging.log(2, 'Defaults of %s : %s', f, f.__defaults__)
+  logging.log(2, 'KW defaults of %s : %s', f, f.__kwdefaults__)
+
+  if kwargs is not None:
+    callargs = tf_inspect.getcallargs(f, *args, **kwargs)
+  else:
+    callargs = tf_inspect.getcallargs(f, *args)
+
+  formatted_callargs = '\n'.join(
+      '    {}: {}'.format(k, v) for k, v in callargs.items())
+  logging.log(2, 'Calling %s with\n%s\n', f, formatted_callargs)
+
+
+#
+# Public API
+#
+
+
+@tf_export('autograph.experimental.do_not_convert')
+def do_not_convert(func=None):
+  """Decorator that suppresses the conversion of a function.
+
+  Args:
+    func: function to decorate.
+
+  Returns:
+    If `func` is not None, returns a `Callable` which is equivalent to
+    `func`, but is not converted by AutoGraph.
+    If `func` is None, returns a decorator that, when invoked with a
+    single `func` argument, returns a `Callable` equivalent to the
+    above case.
+  """
+  if func is None:
+    return do_not_convert
+
+  def wrapper(*args, **kwargs):
+    with ag_ctx.ControlStatusCtx(status=ag_ctx.Status.DISABLED):
+      return func(*args, **kwargs)
+
+  if inspect.isfunction(func) or inspect.ismethod(func):
+    wrapper = functools.update_wrapper(wrapper, func)
+
+  return autograph_artifact(wrapper)
+
+
+# TODO(mdan): Make private.
+def convert(recursive=False,
+            optional_features=None,
+            user_requested=True,
+            conversion_ctx=ag_ctx.NullCtx()):
+  """Decorator that compiles a function to use TensorFlow ops.
+
+  The decorator is dynamic - it recompiles the target whenever the decorated
+  function is called. This means the parameter values are known at conversion.
+  It also means that repeated calls with different types of parameters will be
+  correctly processed.
+
+  Args:
+    recursive: bool, whether to recursively convert any functions or classes
+      that the converted function may use.
+    optional_features: converted.Feature, allows toggling optional or
+      experimental features. When set to None, only the core features are
+      enabled.
+    user_requested: bool, whether this is a function that the user explicitly
+      asked to be converted. See ConversionOptions.user_requested.
+    conversion_ctx: Optional ag_ctx.ControlStatusCtx, the Autograph context in
+      which `f` is used.
+
+  Returns:
+    Callable, a decorator that converts the given function into an equivalent
+    function that uses TensorFlow ops.
+  """
+
+  def decorator(f):
+    """Decorator implementation."""
+
+    def wrapper(*args, **kwargs):
+      """Wrapper that calls the converted version of f."""
+      options = converter.ConversionOptions(
+          recursive=recursive,
+          user_requested=user_requested,
+          optional_features=optional_features)
+      try:
+        with conversion_ctx:
+          return converted_call(f, args, kwargs, options=options)
+      except Exception as e:  # pylint:disable=broad-except
+        if hasattr(e, 'ag_error_metadata'):
+          raise e.ag_error_metadata.to_exception(e)
+        else:
+          raise
+
+    if inspect.isfunction(f) or inspect.ismethod(f):
+      wrapper = functools.update_wrapper(wrapper, f)
+
+    decorated_wrapper = tf_decorator.make_decorator(f, wrapper)
+    return autograph_artifact(decorated_wrapper)
+
+  return decorator
+
+
+# pylint:disable=line-too-long
+@tf_export('autograph.to_graph', v1=[])
+def to_graph(entity, recursive=True, experimental_optional_features=None):
+  """Converts a Python entity into a TensorFlow graph.
+
+  Also see: `tf.autograph.to_code`, `tf.function`.
+
+  Unlike `tf.function`, `to_graph` is a low-level transpiler that converts
+  Python code to TensorFlow graph code. It does not implement any caching,
+  variable management or create any actual ops, and is best used where greater
+  control over the generated TensorFlow graph is desired. Another difference
+  from `tf.function` is that `to_graph` will not wrap the graph into a
+  TensorFlow function or a Python callable. Internally, `tf.function` uses
+  `to_graph`.
+
+  Example usage:
+
+  >>> def f(x):
+  ...   if x > 0:
+  ...     y = x * x
+  ...   else:
+  ...     y = -x
+  ...   return y
+  ...
+  >>> converted_f = to_graph(f)
+  >>> x = tf.constant(2)
+  >>> converted_f(x)  # converted_foo is like a TensorFlow Op.
+  <tf.Tensor: shape=(), dtype=int32, numpy=4>
+
+  Supported Python entities include:
+    * functions
+    * classes
+    * object methods
+
+  Functions are converted into new functions with converted code.
+
+  Classes are converted by generating a new class whose methods use converted
+  code.
+
+  Methods are converted into unbound function that have an additional first
+  argument called `self`.
+
+  For a tutorial, see the
+  [tf.function and AutoGraph guide](https://www.tensorflow.org/guide/function).
+  For more detailed information, see the
+  [AutoGraph reference documentation](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/autograph/g3doc/reference/index.md).
+
+  Args:
+    entity: Python callable or class to convert.
+    recursive: Whether to recursively convert any functions that the converted
+      function may call.
+    experimental_optional_features: `None`, a tuple of, or a single
+      `tf.autograph.experimental.Feature` value.
+
+  Returns:
+    Same as `entity`, the converted Python function or class.
+
+  Raises:
+    ValueError: If the entity could not be converted.
+  """
+  try:
+    program_ctx = converter.ProgramContext(
+        options=converter.ConversionOptions(
+            recursive=recursive,
+            user_requested=True,
+            optional_features=experimental_optional_features))
+    return autograph_artifact(_convert_actual(entity, program_ctx))
+  except (ValueError, AttributeError, KeyError, NameError, AssertionError) as e:
+    logging.error(1, 'Error converting %s', entity, exc_info=True)
+    raise ConversionError('converting {}: {}: {}'.format(
+        entity, e.__class__.__name__, str(e)))
+
+
+@tf_export(v1=['autograph.to_graph'])
+def to_graph_v1(entity,
+                recursive=True,
+                arg_values=None,
+                arg_types=None,
+                experimental_optional_features=None):
+  """Converts a Python entity into a TensorFlow graph.
+
+  Also see: `tf.autograph.to_code`, `tf.function`.
+
+  Unlike `tf.function`, `to_graph` is a low-level transpiler that converts
+  Python code to TensorFlow graph code. It does not implement any caching,
+  variable management or create any actual ops, and is best used where greater
+  control over the generated TensorFlow graph is desired. Another difference
+  from `tf.function` is that `to_graph` will not wrap the graph into a
+  TensorFlow function or a Python callable. Internally, `tf.function` uses
+  `to_graph`.
+
+  _Example Usage_
+
+  ```python
+    def foo(x):
+      if x > 0:
+        y = x * x
+      else:
+        y = -x
+      return y
+
+    converted_foo = to_graph(foo)
+
+    x = tf.constant(1)
+    y = converted_foo(x)  # converted_foo is a TensorFlow Op-like.
+    assert is_tensor(y)
+  ```
+
+  Supported Python entities include:
+    * functions
+    * classes
+    * object methods
+
+  Functions are converted into new functions with converted code.
+
+  Classes are converted by generating a new class whose methods use converted
+  code.
+
+  Methods are converted into unbound function that have an additional first
+  argument called `self`.
+
+  Args:
+    entity: Python callable or class to convert.
+    recursive: Whether to recursively convert any functions that the converted
+      function may call.
+    arg_values: Deprecated.
+    arg_types: Deprecated.
+    experimental_optional_features: `None`, a tuple of, or a single
+      `tf.autograph.experimental.Feature` value.
+
+  Returns:
+    Same as `entity`, the converted Python function or class.
+
+  Raises:
+    ValueError: If the entity could not be converted.
+  """
+  del arg_types
+  del arg_values
+  return to_graph(
+      entity,
+      recursive=recursive,
+      experimental_optional_features=experimental_optional_features)
+
+
+@tf_export(v1=['autograph.to_code'])
+def to_code_v1(entity,
+               recursive=True,
+               arg_values=None,
+               arg_types=None,
+               indentation='  ',
+               experimental_optional_features=None):
+  """Returns the source code generated by AutoGraph, as a string.
+
+  Example usage:
+
+  >>> def f(x):
+  ...   if x < 0:
+  ...     x = -x
+  ...   return x
+  >>> tf.autograph.to_code(f)
+  "...def tf__f(x):..."
+
+  Also see: `tf.autograph.to_graph`.
+
+  Note: If a function has been decorated with `tf.function`, pass its
+  underlying Python function, rather than the callable that `tf.function
+  creates:
+
+  >>> @tf.function
+  ... def f(x):
+  ...   if x < 0:
+  ...     x = -x
+  ...   return x
+  >>> tf.autograph.to_code(f.python_function)
+  "...def tf__f(x):..."
+
+  Args:
+    entity: Python callable or class.
+    recursive: Whether to recursively convert any functions that the converted
+      function may call.
+    arg_values: Deprecated.
+    arg_types: Deprecated.
+    indentation: Deprecated.
+    experimental_optional_features: `None`, a tuple of, or a single
+      `tf.autograph.experimental.Feature` value.
+
+  Returns:
+    The converted code as string.
+  """
+  del arg_values
+  del arg_types
+  del indentation
+  return to_code(
+      entity,
+      recursive=recursive,
+      experimental_optional_features=experimental_optional_features)
+
+
+@tf_export('autograph.to_code', v1=[])
+def to_code(entity, recursive=True, experimental_optional_features=None):
+  """Returns the source code generated by AutoGraph, as a string.
+
+  Example usage:
+
+  >>> def f(x):
+  ...   if x < 0:
+  ...     x = -x
+  ...   return x
+  >>> tf.autograph.to_code(f)
+  "...def tf__f(x):..."
+
+  Also see: `tf.autograph.to_graph`.
+
+  Note: If a function has been decorated with `tf.function`, pass its
+  underlying Python function, rather than the callable that `tf.function
+  creates:
+
+  >>> @tf.function
+  ... def f(x):
+  ...   if x < 0:
+  ...     x = -x
+  ...   return x
+  >>> tf.autograph.to_code(f.python_function)
+  "...def tf__f(x):..."
+
+  Args:
+    entity: Python callable or class to convert.
+    recursive: Whether to recursively convert any functions that the converted
+      function may call.
+    experimental_optional_features: `None`, a tuple of, or a single
+      `tf.autograph.experimental.Feature` value.
+
+  Returns:
+    The converted code as string.
+  """
+  source = tf_inspect.getsource(
+      to_graph(
+          entity,
+          recursive=recursive,
+          experimental_optional_features=experimental_optional_features))
+  return textwrap.dedent(source)
+
+
+_TRANSPILER = PyToTF()

tensorflow/python/autograph/impl/conversion.py

@@ -0,0 +1,227 @@
+# Copyright 2016 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.
+# ==============================================================================
+"""Core conversion logic, serves as main point of access."""
+
+import functools
+import inspect
+import sys
+import unittest
+
+from tensorflow.python.autograph.core import config
+from tensorflow.python.autograph.pyct import cache
+from tensorflow.python.autograph.pyct import inspect_utils
+from tensorflow.python.autograph.utils import ag_logging as logging
+from tensorflow.python.eager.polymorphic_function import tf_method_target
+from tensorflow.python.util import tf_inspect
+
+
+_ALLOWLIST_CACHE = cache.UnboundInstanceCache()
+
+
+def _is_of_known_loaded_module(f, module_name):
+  mod = sys.modules.get(module_name, None)
+  if mod is None:
+    return False
+  if any(v is not None for v in mod.__dict__.values() if f is v):
+    return True
+  return False
+
+
+def _is_known_loaded_type(f, module_name, entity_name):
+  """Tests whether the function or method is an instance of a known type."""
+  if (module_name not in sys.modules or
+      not hasattr(sys.modules[module_name], entity_name)):
+    return False
+  type_entity = getattr(sys.modules[module_name], entity_name)
+  if isinstance(f, type_entity):
+    # The method if of this type. Example:
+    #
+    # o = ClassType()
+    # function(o.method)()
+    return True
+  # Note: inspect is required here, to avoid unpacking tf.function decorators.
+  if inspect.ismethod(f):
+    # The unbound method if of this type. Example:
+    #
+    # class ClassType:
+    #   @function
+    #   def method(self):
+    #     ...
+    # o = ClassType()
+    # o.method()
+    if isinstance(f.__func__, type_entity):
+      return True
+  return False
+
+
+def is_unsupported(o):
+  """Checks whether an entity is supported by AutoGraph at all."""
+
+  # TODO(b/122265385): Remove this bypass.
+  if (_is_known_loaded_type(o, 'wrapt', 'FunctionWrapper') or
+      _is_known_loaded_type(o, 'wrapt', 'BoundFunctionWrapper')):
+    logging.warning(
+        '{} appears to be decorated by wrapt, which is not yet supported'
+        ' by AutoGraph. The function will run as-is.'
+        ' You may still apply AutoGraph before the wrapt decorator.'.format(o))
+    logging.log(2, 'Permanently allowed: %s: wrapt decorated', o)
+    return True
+
+  if _is_known_loaded_type(o, 'functools', '_lru_cache_wrapper'):
+    logging.log(2, 'Permanently allowed: %s: lru_cache', o)
+    return True
+
+  # Constructors are permanently allowed.
+  # TODO(mdan): Toggle as experimental feature instead.
+  # TODO(b/124016764): Remove this limitation.
+  if inspect_utils.isconstructor(o):
+    logging.log(2, 'Permanently allowed: %s: constructor', o)
+    return True
+
+  # Other built-in modules are permanently allowed.
+  # TODO(mdan): Figure out how to do this consistently for all stdlib modules.
+  if any(
+      _is_of_known_loaded_module(o, m)
+      for m in ('collections', 'pdb', 'copy', 'inspect', 're')):
+    logging.log(2, 'Permanently allowed: %s: part of builtin module', o)
+    return True
+
+  # Custom ops and kernels are also permanently allowed.
+  # See tensorflow.framework.load_library.
+  if (hasattr(o, '__module__') and
+      hasattr(o.__module__, '_IS_TENSORFLOW_PLUGIN')):
+    logging.log(2, 'Permanently allowed: %s: TensorFlow plugin', o)
+    return True
+
+  return False
+
+
+# TODO(mdan): allow_namedtuple_subclass should be hardcoded to True.
+def is_allowlisted(
+    o, check_call_override=True, allow_namedtuple_subclass=False):
+  """Checks whether an entity is allowed for use in graph mode.
+
+  Examples of allowed entities include all members of the tensorflow
+  package.
+
+  Args:
+    o: A Python entity.
+    check_call_override: Reserved for internal use. When set to `False`, it
+      disables the rule according to which classes are allowed if their
+      __call__ method is allowed.
+    allow_namedtuple_subclass: Reserved for internal use. When `True`,
+      namedtuple subclasses are not allowed.
+
+  Returns:
+    Boolean
+  """
+  # TODO(b/120224672): Fix this.
+  if isinstance(o, functools.partial):
+    # tf_inspect.getmodule(functools.partial(...)) otherwise returns None since
+    # functools.partial objects do not have a __module__ attribute.
+    m = functools
+  else:
+    m = tf_inspect.getmodule(o)
+
+  # Examples of callables that lack a __module__ property include builtins.
+  if hasattr(m, '__name__'):
+    for rule in config.CONVERSION_RULES:
+      action = rule.get_action(m)
+      if action == config.Action.CONVERT:
+        logging.log(2, 'Not allowed: %s: %s', o, rule)
+        return False
+      elif action == config.Action.DO_NOT_CONVERT:
+        logging.log(2, 'Allowlisted: %s: %s', o, rule)
+        return True
+
+  # The check for __code__ below is because isgeneratorfunction crashes
+  # without one.
+  if hasattr(o, '__code__') and tf_inspect.isgeneratorfunction(o):
+    logging.log(2, 'Allowlisted: %s: generator functions are not converted', o)
+    return True
+
+  if (check_call_override and not tf_inspect.isclass(o) and
+      hasattr(o, '__call__')):
+    # Callable objects: allowed if their __call__ method is.
+    # The type check avoids infinite recursion around the __call__ method
+    # of function objects.
+    if (type(o) != type(o.__call__)) and is_allowlisted(o.__call__):  # pylint: disable=unidiomatic-typecheck
+      logging.log(2, 'Allowlisted: %s: object __call__ allowed', o)
+      return True
+
+  owner_class = None
+  if tf_inspect.ismethod(o):
+    # Methods of allowed classes are also allowed, even if they are
+    # bound via user subclasses.
+    #
+    # For example, suppose `tf.Foo` has a method called `bar`, and `baz` is
+    # defined as below. `tf.Foo` is allowed. Then `baz.bar` is also
+    # allowed.
+    #
+    #   class Custom(tf.Foo):
+    #     pass
+    #
+    #   baz = Custom()
+    #
+    # For the example above, if `Custom` did overload `bar`, then it would no
+    # longer be allowed.
+
+    owner_class = inspect_utils.getmethodclass(o)
+    if owner_class is tf_method_target.TfMethodTarget:
+      owner_class = o.__self__.target_class
+    if owner_class is not None:
+      if issubclass(owner_class, unittest.TestCase):
+        logging.log(2, 'Allowlisted: %s: method of TestCase subclass', o)
+        return True
+
+      owner_class = inspect_utils.getdefiningclass(o, owner_class)
+      if is_allowlisted(
+          owner_class,
+          check_call_override=False,
+          allow_namedtuple_subclass=True):
+        logging.log(2, 'Allowlisted: %s: owner is allowed %s', o,
+                    owner_class)
+        return True
+
+  if inspect_utils.isnamedtuple(o):
+    # Due to the way they're constructed, namedtuple types cannot be converted
+    # because they don't expose source code. But we assume they are safe for
+    # graph mode since they are just containers.
+    if allow_namedtuple_subclass:
+      if not any(inspect_utils.isnamedtuple(base) for base in o.__bases__):
+        logging.log(2, 'Allowlisted: %s: named tuple', o)
+        return True
+    else:
+      logging.log(2, 'Allowlisted: %s: named tuple or subclass', o)
+      return True
+
+  logging.log(2, 'Not allowed: %s: default rule', o)
+  return False
+
+
+def is_in_allowlist_cache(entity, options):
+  try:
+    return _ALLOWLIST_CACHE.has(entity, options)
+  except TypeError:
+    # Catch-all for entities that are unhashable or don't allow weakrefs.
+    return False
+
+
+def cache_allowlisted(entity, options):
+  try:
+    _ALLOWLIST_CACHE[entity][options] = True
+  except TypeError:
+    # Catch-all for entities that are unhashable or don't allow weakrefs.
+    pass

tensorflow/python/autograph/impl/conversion_test.py

@@ -0,0 +1,123 @@
+# 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 conversion module."""
+
+import sys
+import types
+import weakref
+
+from tensorflow.python.autograph import utils
+from tensorflow.python.autograph.core import config
+from tensorflow.python.autograph.core import converter
+from tensorflow.python.autograph.impl import api
+from tensorflow.python.autograph.impl import conversion
+from tensorflow.python.autograph.impl.testing import pybind_for_testing
+from tensorflow.python.eager import function
+from tensorflow.python.framework import constant_op
+from tensorflow.python.platform import test
+
+
+class ConversionTest(test.TestCase):
+
+  def _simple_program_ctx(self):
+    return converter.ProgramContext(
+        options=converter.ConversionOptions(recursive=True),
+        autograph_module=api)
+
+  def test_is_allowlisted(self):
+
+    def test_fn():
+      return constant_op.constant(1)
+
+    self.assertFalse(conversion.is_allowlisted(test_fn))
+    self.assertTrue(conversion.is_allowlisted(utils))
+    self.assertTrue(conversion.is_allowlisted(constant_op.constant))
+
+  def test_is_allowlisted_tensorflow_like(self):
+
+    tf_like = types.ModuleType('tensorflow_foo')
+    def test_fn():
+      pass
+    tf_like.test_fn = test_fn
+    test_fn.__module__ = tf_like
+
+    self.assertFalse(conversion.is_allowlisted(tf_like.test_fn))
+
+  def test_is_allowlisted_callable_allowlisted_call(self):
+
+    allowlisted_mod = types.ModuleType('test_allowlisted_call')
+    sys.modules['test_allowlisted_call'] = allowlisted_mod
+    config.CONVERSION_RULES = ((config.DoNotConvert('test_allowlisted_call'),) +
+                               config.CONVERSION_RULES)
+
+    class TestClass:
+
+      def __call__(self):
+        pass
+
+      def allowlisted_method(self):
+        pass
+
+    TestClass.__module__ = 'test_allowlisted_call'
+    TestClass.__call__.__module__ = 'test_allowlisted_call'
+
+    class Subclass(TestClass):
+
+      def converted_method(self):
+        pass
+
+    tc = Subclass()
+
+    self.assertTrue(conversion.is_allowlisted(TestClass.__call__))
+    self.assertTrue(conversion.is_allowlisted(tc))
+    self.assertTrue(conversion.is_allowlisted(tc.__call__))
+    self.assertTrue(conversion.is_allowlisted(tc.allowlisted_method))
+    self.assertFalse(conversion.is_allowlisted(Subclass))
+    self.assertFalse(conversion.is_allowlisted(tc.converted_method))
+
+  def test_is_allowlisted_tfmethodwrapper(self):
+    allowlisted_mod = types.ModuleType('test_allowlisted_call')
+    sys.modules['test_allowlisted_call'] = allowlisted_mod
+    config.CONVERSION_RULES = ((config.DoNotConvert('test_allowlisted_call'),) +
+                               config.CONVERSION_RULES)
+
+    class TestClass:
+
+      def member_function(self):
+        pass
+
+    TestClass.__module__ = 'test_allowlisted_call'
+    test_obj = TestClass()
+
+    def test_fn(self):
+      del self
+
+    bound_method = types.MethodType(
+        test_fn,
+        function.TfMethodTarget(
+            weakref.ref(test_obj), test_obj.member_function))
+
+    self.assertTrue(conversion.is_allowlisted(bound_method))
+
+  def test_is_allowlisted_pybind(self):
+    test_object = pybind_for_testing.TestClassDef()
+    with test.mock.patch.object(config, 'CONVERSION_RULES', ()):
+      # TODO(mdan): This should return True for functions and methods.
+      # Note: currently, native bindings are allowlisted by a separate check.
+      self.assertFalse(conversion.is_allowlisted(test_object.method))
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/autograph/impl/testing/BUILD

@@ -0,0 +1,19 @@
+load("//tensorflow:tensorflow.default.bzl", "pybind_extension")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//tensorflow:__subpackages__"],
+    licenses = ["notice"],
+)
+
+pybind_extension(
+    name = "pybind_for_testing",
+    srcs = ["pybind_for_testing.cc"],
+    enable_stub_generation = True,
+    pytype_srcs = [
+        "pybind_for_testing.pyi",
+    ],
+    deps = [
+        "@pybind11",
+    ],
+)

tensorflow/python/autograph/impl/testing/pybind_for_testing.cc

@@ -0,0 +1,39 @@
+// 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.
+// ==============================================================================
+
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "pybind11/pytypes.h"  // from @pybind11
+#include "pybind11/stl.h"  // from @pybind11
+
+namespace autograph {
+
+namespace py = pybind11;
+
+class TestClassDef {
+ public:
+  TestClassDef() = default;
+
+  py::object Method() const;
+};
+
+py::object TestClassDef::Method() const { return py::none(); }
+
+PYBIND11_MODULE(pybind_for_testing, m) {
+  py::class_<TestClassDef>(m, "TestClassDef")
+      .def(py::init<>())
+      .def("method", &TestClassDef::Method);
+}
+
+}  // namespace autograph

tensorflow/python/autograph/impl/testing/pybind_for_testing.pyi

@@ -0,0 +1,18 @@
+# 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.
+# ==============================================================================
+
+class TestClassDef:
+    def __init__(self) -> None: ...
+    def method(self) -> object: ...

tensorflow/python/autograph/lang/BUILD

@@ -0,0 +1,45 @@
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:tensorflow.bzl", "py_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "directives",
+    srcs = ["directives.py"],
+    strict_deps = True,
+    visibility = ["//tensorflow:__subpackages__"],
+    deps = [
+        "//tensorflow/python/util:tf_export",
+    ],
+)
+
+py_library(
+    name = "special_functions",
+    srcs = ["special_functions.py"],
+    strict_deps = True,
+    visibility = ["//tensorflow:__subpackages__"],
+    deps = [
+        "//tensorflow/python/autograph/operators:data_structures",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:tensor_util",
+    ],
+)
+
+py_test(
+    name = "special_functions_test",
+    srcs = ["special_functions_test.py"],
+    strict_deps = True,
+    deps = [
+        #internal proto upb dep
+        "//third_party/py/numpy",
+        "//tensorflow/python/autograph/lang:special_functions",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:tensor_util",
+        "//tensorflow/python/ops:list_ops",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)

tensorflow/python/autograph/lang/directives.py

@@ -0,0 +1,94 @@
+# 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.
+# ==============================================================================
+"""Directives are special no-op functions that serve as compilation markers.
+
+They provide static information like type hints, compilation and TensorFlow
+overrides.
+
+These serve as annotations in the compiled code, allowing the user some control
+over the compilation process. They have no functional role at runtime.
+"""
+
+from tensorflow.python.util.tf_export import tf_export
+
+UNSPECIFIED = object()
+
+
+def set_element_type(entity, dtype, shape=UNSPECIFIED):
+  """Indicates that the entity is expected hold items of specified type/shape.
+
+  The staged TensorFlow ops will reflect and assert this data type. Ignored
+  otherwise.
+
+  Args:
+    entity: The entity to annotate.
+    dtype: TensorFlow dtype value to assert for entity.
+    shape: Optional shape to assert for entity.
+  """
+  del entity
+  del dtype
+  del shape
+
+
+@tf_export('autograph.experimental.set_loop_options')
+def set_loop_options(
+    parallel_iterations=UNSPECIFIED,
+    swap_memory=UNSPECIFIED,
+    maximum_iterations=UNSPECIFIED,
+    shape_invariants=UNSPECIFIED):
+  """Specifies additional arguments to be passed to the enclosing while_loop.
+
+  The parameters apply to and only to the immediately enclosing loop. It only
+  has effect if the loop is staged as a TF while_loop; otherwise the parameters
+  have no effect.
+
+  Usage:
+
+    >>> @tf.function(autograph=True)
+    ... def f():
+    ...   n = 0
+    ...   for i in tf.range(10):
+    ...     tf.autograph.experimental.set_loop_options(maximum_iterations=3)
+    ...     n += 1
+    ...   return n
+
+    >>> @tf.function(autograph=True)
+    ... def f():
+    ...   v = tf.constant((0,))
+    ...   for i in tf.range(3):
+    ...     tf.autograph.experimental.set_loop_options(
+    ...         shape_invariants=[(v, tf.TensorShape([None]))]
+    ...     )
+    ...     v = tf.concat((v, [i]), 0)
+    ...   return v
+
+  Also see tf.while_loop.
+
+  Args:
+    parallel_iterations: The maximum number of iterations allowed to run in
+        parallel at any given time. Note that this does not guarantee parallel
+        execution.
+    swap_memory: Whether to store intermediate values needed for
+        gradients on the CPU instead of GPU.
+    maximum_iterations: Allows limiting the total number of iterations executed
+        by the loop.
+    shape_invariants: Allows controlling the argument with the same name passed
+        to tf.while_loop. Unlike tf.while_loop, this is a list of
+        `(tensor, shape)` pairs.
+  """
+  del parallel_iterations
+  del swap_memory
+  del maximum_iterations
+  del shape_invariants

tensorflow/python/autograph/lang/special_functions.py

@@ -0,0 +1,118 @@
+# 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.
+# ==============================================================================
+"""Special functions that only make sense for AutoGraph.
+
+These functions are meant to ensure feature parity between Python and AutoGraph,
+so that the exact same code works in both modes. In general, AutoGraph will
+replace these calls.
+"""
+
+from tensorflow.python.autograph.operators import data_structures
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import tensor_util
+
+
+def _validate_list_constructor(elements, element_dtype, element_shape):
+  """Validates the inputs of tensor_list."""
+  if element_dtype is not None and element_shape is not None:
+    return
+  if tensor_util.is_tf_type(elements):
+    return
+  if isinstance(elements, (list, tuple)):
+    if elements:
+      return
+    else:
+      raise ValueError(
+          'element_dtype and element_shape are required when elements are'
+          ' empty')
+
+  raise ValueError(
+      'unknown type for elements: {}; only Tensor, list and tuple are'
+      ' allowed'.format(type(elements)))
+
+
+def match_staging_level(value, like_value):
+  """Casts a value to be staged at the same level as another."""
+  if tensor_util.is_tf_type(like_value):
+    return constant_op.constant(value)
+  return value
+
+
+def tensor_list(elements,
+                element_dtype=None,
+                element_shape=None,
+                use_tensor_array=False):
+  """Creates an tensor list and populates it with the given elements.
+
+  This function provides a more uniform access to tensor lists and tensor
+  arrays, and allows optional initialization.
+
+  Note: this function is a simplified wrapper. If you need greater control,
+  it is recommended to use the underlying implementation directly.
+
+  Args:
+    elements: Iterable[tf.Tensor, ...], the elements to initially fill the list
+        with
+    element_dtype: Optional[tf.DType], data type for the elements in the list;
+        required if the list is empty
+    element_shape: Optional[tf.TensorShape], shape for the elements in the list;
+        required if the list is empty
+    use_tensor_array: bool, whether to use the more compatible but restrictive
+        tf.TensorArray implementation
+  Returns:
+    Union[tf.Tensor, tf.TensorArray], the new list.
+  Raises:
+    ValueError: for invalid arguments
+  """
+  _validate_list_constructor(elements, element_dtype, element_shape)
+  if use_tensor_array:
+    return data_structures.tf_tensor_array_new(elements, element_dtype,
+                                               element_shape)
+  else:
+    return data_structures.tf_tensor_list_new(elements, element_dtype,
+                                              element_shape)
+
+
+def stack(list_or_tensor, element_dtype=None, strict=True):
+  """Stacks the input, if it admits the notion of stacking.
+
+  For example, a list of tensors can be stacked into a larger tensor. This
+  function is similar to tf.stack, but it accepts non-lists and lists of
+  non-tensors as arguments. In the latter case, the function does nothing.
+
+  Args:
+    list_or_tensor: Any
+    element_dtype: tf.DType, optional dtypedtype for the elements in the list.
+        Required if the input is stackable, and the list is untyped.
+    strict: bool, if True an error is raised if the input is not stackable.
+        Otherwise the function is a no-op.
+
+  Returns:
+    Any, if the input is stackable, the result will be a tf.Tensor. Otherwise,
+    if strict=False, the result will be list_or_tensor.
+
+  Raises:
+    ValueError: if strict=True and the input is not stackable.
+  """
+  if strict:
+    def raise_error(x):
+      raise ValueError('%s must be stackable when strict=True' % x)
+    original_call = raise_error
+  else:
+    original_call = lambda x: x
+  return data_structures.list_stack(
+      list_or_tensor,
+      data_structures.ListStackOpts(
+          element_dtype=element_dtype, original_call=original_call))

tensorflow/python/autograph/lang/special_functions_test.py

@@ -0,0 +1,108 @@
+# 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 special_functions module."""
+
+import numpy as np
+
+from tensorflow.python.autograph.lang import special_functions
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor_util
+from tensorflow.python.ops import list_ops
+from tensorflow.python.platform import test
+
+
+class SpecialFunctionsTest(test.TestCase):
+
+  def test_match_staging_level(self):
+    some_tensor = constant_op.constant(0)
+    tensor_one = special_functions.match_staging_level(1, some_tensor)
+    python_one = special_functions.match_staging_level(1, 1)
+    with self.cached_session() as sess:
+      self.assertTrue(tensor_util.is_tf_type(tensor_one))
+      self.assertAllEqual(self.evaluate(tensor_one), 1)
+      self.assertEqual(python_one, 1)
+
+  def test_tensor_list_empty_list(self):
+    l = special_functions.tensor_list([],
+                                      element_dtype=dtypes.int32,
+                                      element_shape=())
+    sl = list_ops.tensor_list_stack(l, element_dtype=dtypes.int32)
+    with self.cached_session() as sess:
+      self.assertAllEqual(self.evaluate(sl), [])
+
+    l = special_functions.tensor_list((),
+                                      element_dtype=dtypes.int32,
+                                      element_shape=())
+    sl = list_ops.tensor_list_stack(l, element_dtype=dtypes.int32)
+    with self.cached_session() as sess:
+      self.assertAllEqual(self.evaluate(sl), [])
+
+  def test_tensor_list_tensor(self):
+    l = special_functions.tensor_list(
+        constant_op.constant([], dtype=dtypes.int32))
+    sl = list_ops.tensor_list_stack(l, element_dtype=dtypes.int32)
+    with self.cached_session() as sess:
+      self.assertAllEqual(self.evaluate(sl), [])
+
+  def test_tensor_list_unsupported_initializer(self):
+    with self.assertRaisesRegex(ValueError, 'unknown type'):
+      special_functions.tensor_list(np.array([1, 2, 3]))
+
+  def test_tensor_list_empty_list_no_type(self):
+    with self.assertRaisesRegex(ValueError,
+                                'element_dtype and element_shape are required'):
+      special_functions.tensor_list([])
+
+  def test_tensor_list_from_elements(self):
+    elements = [constant_op.constant([1, 2]), constant_op.constant([3, 4])]
+
+    l = special_functions.tensor_list(elements)
+    sl = list_ops.tensor_list_stack(l, element_dtype=dtypes.int32)
+    with self.cached_session() as sess:
+      self.assertAllEqual(self.evaluate(sl), [[1, 2], [3, 4]])
+
+  def test_tensor_list_array_from_elements(self):
+    elements = [constant_op.constant([1, 2]), constant_op.constant([3, 4])]
+
+    l = special_functions.tensor_list(elements, use_tensor_array=True)
+    sl = l.stack()
+    with self.cached_session() as sess:
+      self.assertAllEqual(self.evaluate(sl), [[1, 2], [3, 4]])
+
+  def test_stack(self):
+    self.assertEqual(special_functions.stack(1, strict=False), 1)
+    self.assertListEqual(
+        special_functions.stack([1, 2, 3], strict=False), [1, 2, 3])
+    # TODO(mdan): This should probably forward to tf.stack.
+    self.assertTrue(
+        isinstance(
+            special_functions.stack(
+                [constant_op.constant(1),
+                 constant_op.constant(2)], strict=False), list))
+
+    with self.assertRaises(ValueError):
+      special_functions.stack([1, 2, 3])
+
+    t = constant_op.constant([1.0, 2.0])
+    l = list_ops.tensor_list_from_tensor(
+        t, element_shape=constant_op.constant([], dtype=dtypes.int32))
+    self.assertTrue(
+        tensor_util.is_tf_type(
+            special_functions.stack(l, element_dtype=dtypes.float32)))
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/autograph/operators/BUILD

@@ -0,0 +1,295 @@
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:pytype.default.bzl", "pytype_strict_library")
+load("//tensorflow:tensorflow.bzl", "py_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "py_builtins",
+    srcs = ["py_builtins.py"],
+    strict_deps = True,
+    visibility = ["//tensorflow:__subpackages__"],
+    deps = [
+        "//tensorflow/python/autograph/utils:tensors",
+        "//tensorflow/python/autograph/utils:type_registry",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor_util",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:cond",
+        "//tensorflow/python/ops:control_flow_assert",
+        "//tensorflow/python/ops:list_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:parsing_ops_gen",
+        "//tensorflow/python/ops:string_ops_gen",
+    ],
+)
+
+py_library(
+    name = "exceptions",
+    srcs = ["exceptions.py"],
+    strict_deps = True,
+    visibility = ["//visibility:private"],
+    deps = [
+        "//tensorflow/python/framework:tensor_util",
+        "//tensorflow/python/ops:control_flow_assert",
+        "//tensorflow/python/util:tf_inspect",
+    ],
+)
+
+py_library(
+    name = "__init__",
+    srcs = ["__init__.py"],
+    strict_deps = True,
+    visibility = ["//tensorflow:__subpackages__"],
+    deps = [
+        ":conditional_expressions",
+        ":control_flow",
+        ":data_structures",
+        ":exceptions",
+        ":logical",
+        ":py_builtins",
+        ":slices",
+        ":variables",
+    ],
+)
+
+py_library(
+    name = "logical",
+    srcs = ["logical.py"],
+    strict_deps = True,
+    visibility = ["//visibility:private"],
+    deps = [
+        "//tensorflow/python/framework:tensor_util",
+        "//tensorflow/python/ops:cond",
+        "//tensorflow/python/ops:math_ops_gen",
+    ],
+)
+
+py_library(
+    name = "variables",
+    srcs = ["variables.py"],
+    strict_deps = True,
+    visibility = ["//tensorflow:__subpackages__"],
+)
+
+py_library(
+    name = "data_structures",
+    srcs = ["data_structures.py"],
+    strict_deps = True,
+    visibility = ["//tensorflow:__subpackages__"],
+    deps = [
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor_util",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:cond",
+        "//tensorflow/python/ops:list_ops",
+        "//tensorflow/python/ops:tensor_array_ops",
+    ],
+)
+
+py_library(
+    name = "conditional_expressions",
+    srcs = ["conditional_expressions.py"],
+    strict_deps = True,
+    visibility = ["//visibility:private"],
+    deps = [
+        ":control_flow",
+        "//tensorflow/python/autograph/utils:tensors",
+        "//tensorflow/python/ops:cond",
+    ],
+)
+
+py_library(
+    name = "control_flow",
+    srcs = ["control_flow.py"],
+    strict_deps = True,
+    visibility = ["//tensorflow:__subpackages__"],
+    deps = [
+        ":py_builtins",
+        ":variables",
+        "//tensorflow/python/autograph/utils:ag_logging",
+        "//tensorflow/python/autograph/utils:misc",
+        "//tensorflow/python/autograph/utils:tensors",
+        "//tensorflow/python/autograph/utils:type_registry",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:func_graph",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor_conversion",
+        "//tensorflow/python/framework:tensor_shape",
+        "//tensorflow/python/framework:tensor_util",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:cond",
+        "//tensorflow/python/ops:control_flow_assert",
+        "//tensorflow/python/ops:control_flow_util",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:tensor_array_ops",
+        "//tensorflow/python/ops:while_loop",
+        "//tensorflow/python/types:distribute",
+        "//tensorflow/python/util:nest",
+        "//tensorflow/python/util:variable_utils",
+        "//third_party/py/numpy",
+    ],
+)
+
+py_library(
+    name = "slices",
+    srcs = ["slices.py"],
+    strict_deps = True,
+    visibility = ["//visibility:private"],
+    deps = [
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:tensor_util",
+        "//tensorflow/python/ops:array_ops_gen",
+        "//tensorflow/python/ops:list_ops",
+        "//tensorflow/python/ops:string_ops_gen",
+        "//tensorflow/python/ops:tensor_array_ops",
+    ],
+)
+
+py_test(
+    name = "data_structures_test",
+    srcs = ["data_structures_test.py"],
+    strict_deps = True,
+    deps = [
+        ":data_structures",
+        #internal proto upb dep
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:tensor",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:list_ops",
+        "//tensorflow/python/ops:tensor_array_ops",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_test(
+    name = "conditional_expressions_test",
+    srcs = ["conditional_expressions_test.py"],
+    strict_deps = True,
+    deps = [
+        ":conditional_expressions",
+        #internal proto upb dep
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_test(
+    name = "control_flow_test",
+    srcs = ["control_flow_test.py"],
+    strict_deps = True,
+    tags = [
+        "no_gpu",  # b/127001953
+    ],
+    deps = [
+        ":control_flow",
+        ":variables",
+        #internal proto upb dep
+        "//third_party/py/numpy",
+        "//tensorflow/python/autograph/utils:ag_logging",
+        "//tensorflow/python/autograph/utils:testing",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:tensor_shape",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:control_flow_assert",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:math_ops_gen",
+        "//tensorflow/python/ops:random_ops",
+        "//tensorflow/python/ops/ragged:ragged_factory_ops",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_test(
+    name = "exceptions_test",
+    srcs = ["exceptions_test.py"],
+    strict_deps = True,
+    deps = [
+        ":exceptions",
+        #internal proto upb dep
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_test(
+    name = "logical_test",
+    srcs = ["logical_test.py"],
+    strict_deps = True,
+    deps = [
+        ":logical",
+        #internal proto upb dep
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_test(
+    name = "py_builtins_test",
+    srcs = ["py_builtins_test.py"],
+    strict_deps = True,
+    deps = [
+        ":data_structures",
+        ":py_builtins",
+        #internal proto upb dep
+        "//tensorflow/python/autograph/core:converter",
+        "//tensorflow/python/autograph/core:function_wrappers",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:dtypes",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:array_ops",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:tensor_array_ops",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_test(
+    name = "slices_test",
+    srcs = ["slices_test.py"],
+    strict_deps = True,
+    deps = [
+        ":slices",
+        #internal proto upb dep
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/ops:list_ops",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_test(
+    name = "variables_test",
+    srcs = ["variables_test.py"],
+    strict_deps = True,
+    deps = [
+        ":variables",
+        #internal proto upb dep
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+pytype_strict_library(
+    name = "dispatch_context",
+    srcs = ["dispatch_context.py"],
+)

tensorflow/python/autograph/operators/__init__.py

@@ -0,0 +1,63 @@
+# 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.
+# ==============================================================================
+"""This module implements operators that AutoGraph overloads.
+
+Note that "operator" is used loosely here, and includes control structures like
+conditionals and loops, implemented in functional form, using for example
+closures for the body.
+"""
+
+# Naming conventions:
+#  * operator names match the name usually used for the respective Python
+#    idiom; examples: for_stmt, list_append
+#  * operator arguments match either of:
+#    - the corresponding Python AST attribute (e.g. the condition of an if
+#      statement is called test) if the operator represents an AST construct
+#    - the names used in the Python docs, if the operator is a function (e.g.
+#      list_ and x for append, see
+#      https://docs.python.org/3.7/tutorial/datastructures.html)
+#
+# All operators may accept a final argument named "opts", of a type that
+# subclasses namedtuple and contains any arguments that are only required
+# for some specializations of the operator.
+
+from tensorflow.python.autograph.operators.conditional_expressions import if_exp
+from tensorflow.python.autograph.operators.control_flow import for_stmt
+from tensorflow.python.autograph.operators.control_flow import if_stmt
+from tensorflow.python.autograph.operators.control_flow import while_stmt
+from tensorflow.python.autograph.operators.data_structures import list_append
+from tensorflow.python.autograph.operators.data_structures import list_pop
+from tensorflow.python.autograph.operators.data_structures import list_stack
+from tensorflow.python.autograph.operators.data_structures import ListPopOpts
+from tensorflow.python.autograph.operators.data_structures import ListStackOpts
+from tensorflow.python.autograph.operators.data_structures import new_list
+from tensorflow.python.autograph.operators.exceptions import assert_stmt
+from tensorflow.python.autograph.operators.logical import and_
+from tensorflow.python.autograph.operators.logical import eq
+from tensorflow.python.autograph.operators.logical import not_
+from tensorflow.python.autograph.operators.logical import not_eq
+from tensorflow.python.autograph.operators.logical import or_
+from tensorflow.python.autograph.operators.py_builtins import float_
+from tensorflow.python.autograph.operators.py_builtins import int_
+from tensorflow.python.autograph.operators.py_builtins import len_
+from tensorflow.python.autograph.operators.py_builtins import print_
+from tensorflow.python.autograph.operators.py_builtins import range_
+from tensorflow.python.autograph.operators.slices import get_item
+from tensorflow.python.autograph.operators.slices import GetItemOpts
+from tensorflow.python.autograph.operators.slices import set_item
+from tensorflow.python.autograph.operators.variables import ld
+from tensorflow.python.autograph.operators.variables import ldu
+from tensorflow.python.autograph.operators.variables import Undefined
+from tensorflow.python.autograph.operators.variables import UndefinedReturnValue

tensorflow/python/autograph/operators/conditional_expressions.py

@@ -0,0 +1,52 @@
+# Copyright 2020 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.
+# ==============================================================================
+"""Conditional expressions (e.g. the ternary if statement)."""
+
+
+from tensorflow.python.autograph.operators import control_flow
+from tensorflow.python.autograph.utils import tensors
+from tensorflow.python.ops import cond as tf_cond
+
+
+def if_exp(cond, if_true, if_false, expr_repr):
+  if tensors.is_dense_tensor(cond):
+    return _tf_if_exp(cond, if_true, if_false, expr_repr)
+  else:
+    return _py_if_exp(cond, if_true, if_false)
+
+
+def _tf_if_exp(cond, if_true, if_false, expr_repr):
+  """Overload of if_exp that stages a TF cond."""
+  # TODO(mdan): Use nonlocal once we no longer need to support py2.
+  true_val = []
+  false_val = []
+
+  def true_fn():
+    true_val.append(if_true())
+    if true_val and false_val:
+      control_flow.verify_single_cond_var(expr_repr, true_val[0], false_val[0])
+    return true_val[0]
+
+  def false_fn():
+    false_val.append(if_false())
+    if true_val and false_val:
+      control_flow.verify_single_cond_var(expr_repr, true_val[0], false_val[0])
+    return false_val[0]
+
+  return tf_cond.cond(cond, true_fn, false_fn)
+
+
+def _py_if_exp(cond, if_true, if_false):
+  return if_true() if cond else if_false()

tensorflow/python/autograph/operators/conditional_expressions_test.py

@@ -0,0 +1,61 @@
+# Copyright 2020 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 conditional_expressions module."""
+
+from tensorflow.python.autograph.operators import conditional_expressions
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import test_util
+from tensorflow.python.platform import test
+
+
+def _basic_expr(cond):
+  return conditional_expressions.if_exp(
+      cond,
+      lambda: constant_op.constant(1),
+      lambda: constant_op.constant(2),
+      'cond')
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class IfExpTest(test.TestCase):
+
+  def test_tensor(self):
+    self.assertEqual(self.evaluate(_basic_expr(constant_op.constant(True))), 1)
+    self.assertEqual(self.evaluate(_basic_expr(constant_op.constant(False))), 2)
+
+  def test_tensor_mismatched_type(self):
+    # tf.function required because eager cond degenerates to Python if.
+    @def_function.function
+    def test_fn():
+      conditional_expressions.if_exp(
+          constant_op.constant(True), lambda: 1.0, lambda: 2, 'expr_repr')
+
+    with self.assertRaisesRegex(
+        TypeError,
+        "'expr_repr' has dtype float32 in the main.*int32 in the else"):
+      test_fn()
+
+  def test_python(self):
+    self.assertEqual(self.evaluate(_basic_expr(True)), 1)
+    self.assertEqual(self.evaluate(_basic_expr(False)), 2)
+    self.assertEqual(
+        conditional_expressions.if_exp(True, lambda: 1, lambda: 2, ''), 1)
+    self.assertEqual(
+        conditional_expressions.if_exp(False, lambda: 1, lambda: 2, ''), 2)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/autograph/operators/data_structures.py

@@ -0,0 +1,347 @@
+# Copyright 2016 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.
+# ==============================================================================
+"""Operators specific to data structures: list append, subscripts, etc."""
+
+import collections
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import cond
+from tensorflow.python.ops import list_ops
+from tensorflow.python.ops import tensor_array_ops
+
+
+# TODO(mdan): Once control flow supports objects, repackage as a class.
+
+
+def new_list(iterable=None):
+  """The list constructor.
+
+  Args:
+    iterable: Optional elements to fill the list with.
+
+  Returns:
+    A list-like object. The exact return value depends on the initial elements.
+  """
+  if iterable:
+    elements = tuple(iterable)
+  else:
+    elements = ()
+
+  if elements:
+    # When the list contains elements, it is assumed to be a "Python" lvalue
+    # list.
+    return _py_list_new(elements)
+  return tf_tensor_list_new(elements)
+
+
+def tf_tensor_array_new(elements, element_dtype=None, element_shape=None):
+  """Overload of new_list that stages a Tensor list creation."""
+  elements = tuple(ops.convert_to_tensor(el) for el in elements)
+
+  all_dtypes = set(el.dtype for el in elements)
+  if len(all_dtypes) == 1:
+    inferred_dtype, = tuple(all_dtypes)
+    if element_dtype is not None and element_dtype != inferred_dtype:
+      raise ValueError(
+          'incompatible dtype; specified: {}, inferred from {}: {}'.format(
+              element_dtype, elements, inferred_dtype))
+  elif len(all_dtypes) > 1:
+    raise ValueError(
+        'TensorArray requires all elements to have the same dtype:'
+        ' {}'.format(elements))
+  else:
+    if element_dtype is None:
+      raise ValueError('dtype is required to create an empty TensorArray')
+
+  all_shapes = set(tuple(el.shape.as_list()) for el in elements)
+  if len(all_shapes) == 1:
+    inferred_shape, = tuple(all_shapes)
+    if element_shape is not None and element_shape != inferred_shape:
+      raise ValueError(
+          'incompatible shape; specified: {}, inferred from {}: {}'.format(
+              element_shape, elements, inferred_shape))
+  elif len(all_shapes) > 1:
+    raise ValueError(
+        'TensorArray requires all elements to have the same shape:'
+        ' {}'.format(elements))
+    # TODO(mdan): We may want to allow different shapes with infer_shape=False.
+  else:
+    inferred_shape = None
+
+  if element_dtype is None:
+    element_dtype = inferred_dtype
+  if element_shape is None:
+    element_shape = inferred_shape
+
+  l = tensor_array_ops.TensorArray(
+      dtype=element_dtype,
+      size=len(elements),
+      dynamic_size=True,
+      infer_shape=(element_shape is None),
+      element_shape=element_shape)
+  for i, el in enumerate(elements):
+    l = l.write(i, el)
+  return l
+
+
+def tf_tensor_list_new(elements, element_dtype=None, element_shape=None):
+  """Overload of new_list that stages a Tensor list creation."""
+  if tensor_util.is_tf_type(elements):
+    if element_shape is not None:
+      raise ValueError(
+          'element shape may not be specified when creating list from tensor')
+    element_shape = array_ops.shape(elements)[1:]
+    l = list_ops.tensor_list_from_tensor(elements, element_shape=element_shape)
+    return l
+
+  elements = tuple(ops.convert_to_tensor(el) for el in elements)
+
+  all_dtypes = set(el.dtype for el in elements)
+  if len(all_dtypes) == 1:
+    inferred_dtype = tuple(all_dtypes)[0]
+    if element_dtype is not None and element_dtype != inferred_dtype:
+      raise ValueError(
+          'incompatible dtype; specified: {}, inferred from {}: {}'.format(
+              element_dtype, elements, inferred_dtype))
+  elif all_dtypes:
+    # Heterogeneous lists are ok.
+    if element_dtype is not None:
+      raise ValueError(
+          'specified dtype {} is inconsistent with that of elements {}'.format(
+              element_dtype, elements))
+    inferred_dtype = dtypes.variant
+  else:
+    inferred_dtype = dtypes.variant
+
+  all_shapes = set(tuple(el.shape.as_list()) for el in elements)
+  if len(all_shapes) == 1:
+    inferred_shape = array_ops.shape(elements[0])
+    if element_shape is not None and element_shape != inferred_shape:
+      raise ValueError(
+          'incompatible shape; specified: {}, inferred from {}: {}'.format(
+              element_shape, elements, inferred_shape))
+  elif all_shapes:
+    # Heterogeneous lists are ok.
+    if element_shape is not None:
+      raise ValueError(
+          'specified shape {} is inconsistent with that of elements {}'.format(
+              element_shape, elements))
+    inferred_shape = constant_op.constant(-1)  # unknown shape, by convention
+  else:
+    inferred_shape = constant_op.constant(-1)  # unknown shape, by convention
+
+  if element_dtype is None:
+    element_dtype = inferred_dtype
+  if element_shape is None:
+    element_shape = inferred_shape
+
+  element_shape = ops.convert_to_tensor(element_shape, dtype=dtypes.int32)
+  l = list_ops.empty_tensor_list(
+      element_shape=element_shape, element_dtype=element_dtype)
+  for el in elements:
+    l = list_ops.tensor_list_push_back(l, el)
+  return l
+
+
+def _py_list_new(elements):
+  """Overload of new_list that creates a Python list."""
+  return list(elements)
+
+
+def list_append(list_, x):
+  """The list append function.
+
+  Note: it is unspecified where list_ will be mutated or not. If list_ is
+  a TensorFlow entity, it will not be typically mutated. If list_ is a plain
+  list, it will be. In general, if the list is mutated then the return value
+  should point to the original entity.
+
+  Args:
+    list_: An entity that supports append semantics.
+    x: The element to append.
+
+  Returns:
+    Same as list_, after the append was performed.
+
+  Raises:
+    ValueError: if list_ is not of a known list-like type.
+  """
+  if isinstance(list_, tensor_array_ops.TensorArray):
+    return _tf_tensorarray_append(list_, x)
+  elif tensor_util.is_tf_type(list_):
+    if list_.dtype == dtypes.variant:
+      return _tf_tensor_list_append(list_, x)
+    else:
+      raise ValueError(
+          'tensor lists are expected to be Tensors with dtype=tf.variant,'
+          ' instead found %s' % list_)
+  else:
+    return _py_list_append(list_, x)
+
+
+def _tf_tensor_list_append(list_, x):
+  """Overload of list_append that stages a Tensor list write."""
+  def empty_list_of_elements_like_x():
+    tensor_x = ops.convert_to_tensor(x)
+    return list_ops.empty_tensor_list(
+        element_shape=array_ops.shape(tensor_x),
+        element_dtype=tensor_x.dtype)
+
+  list_ = cond.cond(
+      list_ops.tensor_list_length(list_) > 0,
+      lambda: list_,
+      empty_list_of_elements_like_x,
+  )
+  return list_ops.tensor_list_push_back(list_, x)
+
+
+def _tf_tensorarray_append(list_, x):
+  """Overload of list_append that stages a TensorArray write."""
+  return list_.write(list_.size(), x)
+
+
+def _py_list_append(list_, x):
+  """Overload of list_append that executes a Python list append."""
+  # Revert to the original call.
+  list_.append(x)
+  return list_
+
+
+class ListPopOpts(
+    collections.namedtuple('ListPopOpts', ('element_dtype', 'element_shape'))):
+  pass
+
+
+def list_pop(list_, i, opts):
+  """The list pop function.
+
+  Note: it is unspecified where list_ will be mutated or not. If list_ is
+  a TensorFlow entity, it will not be typically mutated. If list_ is a plain
+  list, it will be. In general, if the list is mutated then the return value
+  should point to the original entity.
+
+  Args:
+    list_: An entity that supports pop semantics.
+    i: Optional index to pop from. May be None.
+    opts: A ListPopOpts.
+
+  Returns:
+    Tuple (x, out_list_):
+      out_list_: same as list_, after the removal was performed.
+      x: the removed element value.
+
+  Raises:
+    ValueError: if list_ is not of a known list-like type or the operation is
+    not supported for that type.
+  """
+  assert isinstance(opts, ListPopOpts)
+
+  if isinstance(list_, tensor_array_ops.TensorArray):
+    raise ValueError('TensorArray does not support item removal')
+  elif tensor_util.is_tf_type(list_):
+    if list_.dtype == dtypes.variant:
+      return _tf_tensor_list_pop(list_, i, opts)
+    else:
+      raise ValueError(
+          'tensor lists are expected to be Tensors with dtype=tf.variant,'
+          ' instead found %s' % list_)
+  else:
+    return _py_list_pop(list_, i)
+
+
+def _tf_tensor_list_pop(list_, i, opts):
+  """Overload of list_pop that stages a Tensor list pop."""
+  if i is not None:
+    raise NotImplementedError('tensor lists only support removing from the end')
+
+  if opts.element_dtype is None:
+    raise ValueError('cannot pop from a list without knowing its element '
+                     'type; use set_element_type to annotate it')
+  if opts.element_shape is None:
+    raise ValueError('cannot pop from a list without knowing its element '
+                     'shape; use set_element_type to annotate it')
+  list_out, x = list_ops.tensor_list_pop_back(
+      list_, element_dtype=opts.element_dtype)
+  x.set_shape(opts.element_shape)
+  return list_out, x
+
+
+def _py_list_pop(list_, i):
+  """Overload of list_pop that executes a Python list append."""
+  if i is None:
+    x = list_.pop()
+  else:
+    x = list_.pop(i)
+  return list_, x
+
+
+# TODO(mdan): Look into reducing duplication between all these containers.
+class ListStackOpts(
+    collections.namedtuple('ListStackOpts',
+                           ('element_dtype', 'original_call'))):
+  pass
+
+
+def list_stack(list_, opts):
+  """The list stack function.
+
+  This does not have a direct correspondent in Python. The closest idiom to
+  this is tf.append or np.stack. It's different from those in the sense that it
+  accepts a Tensor list, rather than a list of tensors. It can also accept
+  TensorArray. When the target is anything else, the dispatcher will rely on
+  ctx.original_call for fallback.
+
+  Args:
+    list_: An entity that supports append semantics.
+    opts: A ListStackOpts object.
+
+  Returns:
+    The output of the stack operation, typically a Tensor.
+  """
+  assert isinstance(opts, ListStackOpts)
+
+  if isinstance(list_, tensor_array_ops.TensorArray):
+    return _tf_tensorarray_stack(list_)
+  elif tensor_util.is_tf_type(list_):
+    if list_.dtype == dtypes.variant:
+      return _tf_tensor_list_stack(list_, opts)
+    else:
+      # No-op for primitive Tensor arguments.
+      return list_
+  else:
+    return _py_list_stack(list_, opts)
+
+
+def _tf_tensorarray_stack(list_):
+  """Overload of list_stack that stages a TensorArray stack."""
+  return list_.stack()
+
+
+def _tf_tensor_list_stack(list_, opts):
+  """Overload of list_stack that stages a Tensor list write."""
+  if opts.element_dtype is None:
+    raise ValueError('cannot stack a list without knowing its element type;'
+                     ' use set_element_type to annotate it')
+  return list_ops.tensor_list_stack(list_, element_dtype=opts.element_dtype)
+
+
+def _py_list_stack(list_, opts):
+  """Overload of list_stack that executes a Python list append."""
+  # Revert to the original call.
+  return opts.original_call(list_)

tensorflow/python/autograph/operators/data_structures_test.py

@@ -0,0 +1,183 @@
+# 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 data_structures module."""
+
+from tensorflow.python.autograph.operators import data_structures
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import list_ops
+from tensorflow.python.ops import tensor_array_ops
+from tensorflow.python.platform import test
+
+
+class ListTest(test.TestCase):
+
+  def test_new_list_empty(self):
+    l = data_structures.new_list()
+    # Can't evaluate an empty list.
+    # TODO(mdan): sess.run should allow tf.variant maybe?
+    self.assertTrue(isinstance(l, tensor.Tensor))
+
+  def test_new_list_tensor(self):
+    l = data_structures.new_list([3, 4, 5])
+    self.assertAllEqual(l, [3, 4, 5])
+
+  def test_tf_tensor_list_new(self):
+    l = data_structures.tf_tensor_list_new([3, 4, 5])
+    t = list_ops.tensor_list_stack(l, element_dtype=dtypes.int32)
+    with self.cached_session() as sess:
+      self.assertAllEqual(self.evaluate(t), [3, 4, 5])
+
+  def test_tf_tensor_list_new_empty(self):
+    l = data_structures.tf_tensor_list_new([],
+                                           element_dtype=dtypes.int32,
+                                           element_shape=())
+    t = list_ops.tensor_list_stack(l, element_dtype=dtypes.int32)
+    with self.cached_session() as sess:
+      self.assertAllEqual(self.evaluate(t), [])
+
+  def test_tf_tensor_list_new_from_tensor(self):
+    l = data_structures.tf_tensor_list_new(constant_op.constant([3, 4, 5]))
+    t = list_ops.tensor_list_stack(l, element_dtype=dtypes.int32)
+    with self.cached_session() as sess:
+      self.assertAllEqual(self.evaluate(t), [3, 4, 5])
+
+  @test_util.run_deprecated_v1
+  def test_tf_tensor_list_new_illegal_input(self):
+    with self.assertRaises(ValueError):
+      data_structures.tf_tensor_list_new([3, 4.0])
+    # TODO(mdan): It might make more sense to type cast in this case.
+    with self.assertRaises(ValueError):
+      data_structures.tf_tensor_list_new([3, 4], element_dtype=dtypes.float32)
+    # Tensor lists do support heterogeneous lists.
+    self.assertIsNot(data_structures.tf_tensor_list_new([3, [4, 5]]), None)
+    with self.assertRaises(ValueError):
+      data_structures.tf_tensor_list_new([3, 4], element_shape=(2,))
+    with self.assertRaises(ValueError):
+      data_structures.tf_tensor_list_new(
+          constant_op.constant([1, 2, 3]), element_shape=[1])
+
+  def test_tf_tensor_array_new(self):
+    l = data_structures.tf_tensor_array_new([3, 4, 5])
+    t = l.stack()
+    with self.cached_session() as sess:
+      self.assertAllEqual(self.evaluate(t), [3, 4, 5])
+
+  def test_tf_tensor_array_new_illegal_input(self):
+    with self.assertRaises(ValueError):
+      data_structures.tf_tensor_array_new([3, 4.0])
+    with self.assertRaises(ValueError):
+      data_structures.tf_tensor_array_new([3, 4], element_dtype=dtypes.float32)
+    with self.assertRaises(ValueError):
+      data_structures.tf_tensor_array_new([3, [4, 5]])
+    with self.assertRaises(ValueError):
+      data_structures.tf_tensor_array_new([3, 4], element_shape=(2,))
+    with self.assertRaises(ValueError):
+      data_structures.tf_tensor_array_new([], element_shape=(2,))
+    # TAs can infer the shape.
+    self.assertIsNot(
+        data_structures.tf_tensor_array_new([], element_dtype=dtypes.float32),
+        None)
+
+  def test_append_tensor_list(self):
+    l = data_structures.new_list()
+    x = constant_op.constant([1, 2, 3])
+    l = data_structures.list_append(l, x)
+
+    t = list_ops.tensor_list_stack(l, element_dtype=x.dtype)
+    with self.cached_session() as sess:
+      self.assertAllEqual(self.evaluate(t), [[1, 2, 3]])
+
+  @test_util.run_deprecated_v1
+  def test_append_tensorarray(self):
+    l = tensor_array_ops.TensorArray(dtypes.int32, size=0, dynamic_size=True)
+    l1 = data_structures.list_append(l, 1)
+    l2 = data_structures.list_append(l1, 2)
+
+    with self.cached_session() as sess:
+      self.assertAllEqual(self.evaluate(l1.stack()), [1])
+      self.assertAllEqual(self.evaluate(l2.stack()), [1, 2])
+
+  def test_append_python(self):
+    l = []
+    self.assertAllEqual(data_structures.list_append(l, 1), [1])
+    self.assertAllEqual(data_structures.list_append(l, 2), [1, 2])
+
+  def test_pop_tensor_list(self):
+    initial_list = constant_op.constant([[1, 2], [3, 4]])
+    elem_shape = constant_op.constant([2])
+    l = list_ops.tensor_list_from_tensor(initial_list, element_shape=elem_shape)
+
+    opts = data_structures.ListPopOpts(
+        element_dtype=initial_list.dtype,
+        element_shape=(2,))
+
+    with self.assertRaises(NotImplementedError):
+      data_structures.list_pop(l, 0, opts)
+
+    with self.cached_session() as sess:
+      l, x = data_structures.list_pop(l, None, opts)
+      self.assertAllEqual(self.evaluate(x), [3, 4])
+
+      t = list_ops.tensor_list_stack(l, element_dtype=initial_list.dtype)
+      self.assertAllEqual(self.evaluate(t), [[1, 2]])
+
+  def test_pop_python(self):
+    l = [1, 2, 3]
+    opts = data_structures.ListPopOpts(element_dtype=None, element_shape=())
+    self.assertAllEqual(data_structures.list_pop(l, None, opts), ([1, 2], 3))
+    self.assertAllEqual(data_structures.list_pop(l, None, opts), ([1], 2))
+
+  def test_stack_tensor_list(self):
+    initial_list = constant_op.constant([[1, 2], [3, 4]])
+    elem_shape = constant_op.constant([2])
+    l = list_ops.tensor_list_from_tensor(initial_list, element_shape=elem_shape)
+
+    opts = data_structures.ListStackOpts(
+        element_dtype=initial_list.dtype, original_call=None)
+
+    with self.cached_session() as sess:
+      t = data_structures.list_stack(l, opts)
+      self.assertAllEqual(self.evaluate(t), self.evaluate(initial_list))
+
+  @test_util.run_deprecated_v1
+  def test_stack_tensor_list_empty(self):
+    l = list_ops.empty_tensor_list(
+        element_shape=None, element_dtype=dtypes.variant)
+
+    opts = data_structures.ListStackOpts(
+        element_dtype=dtypes.int32, original_call=None)
+
+    # TODO(mdan): Allow stacking empty lists if the dtype and shape are known.
+    with self.assertRaises(ValueError):
+      data_structures.list_stack(l, opts)
+
+  def test_stack_fallback(self):
+
+    def dummy_function(l):
+      # Lazy person's mock: just transform the argument in a way in which we
+      # can check that this function was indeed called.
+      return [x * 2 for x in l]
+
+    opts = data_structures.ListStackOpts(
+        element_dtype=None, original_call=dummy_function)
+
+    self.assertAllEqual(data_structures.list_stack([1, 2], opts), [2, 4])
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/autograph/operators/dispatch_context.py

@@ -0,0 +1,37 @@
+# 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.
+# ==============================================================================
+"""Structures that allow uniform control over the dispatch process."""
+
+import collections
+
+
+# TODO(mdan): This is where macro override controls fit.
+
+
+class DispatchContext(collections.namedtuple(
+    'DispatchContext',
+    ('options',))):
+  """Allows passing additional parameters to the specific implementations.
+
+  Attributes:
+    options: Optional dict of extra arguments that may be required by specific
+      implementations.
+  """
+
+  def option(self, name):
+    return self.options[name]
+
+
+NO_CTX = DispatchContext(options={})

tensorflow/python/autograph/operators/exceptions.py

@@ -0,0 +1,82 @@
+# 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.
+# ==============================================================================
+"""Exception handling statements: assert, etc."""
+
+from tensorflow.python.framework import tensor_util
+from tensorflow.python.ops import control_flow_assert
+from tensorflow.python.util import tf_inspect
+
+
+def assert_stmt(expression1, expression2):
+  """Functional form of an assert statement.
+
+  This follows the semantics of the Python assert statement, however the
+  concrete implementations may deviate from it. See the respective
+  implementation for details.
+
+  In general, the assert statement should not be used for control flow.
+  Furthermore, it is encouraged that the assertion expressions should not have
+  side effects.
+
+  Args:
+    expression1: Any
+    expression2: Callable[[], Any], returns the expression to include in the
+        error message when expression1 evaluates to False. When expression1 is
+        True, the result of expression2 will not be evaluated, however,
+        expression2 itself may be evaluated in some implementations.
+
+  Returns:
+    Any, implementation-dependent.
+
+  Raises:
+    ValueError: if any arguments are illegal.
+  """
+  if not callable(expression2):
+    raise ValueError('{} must be a callable'.format(expression2))
+  args, _, keywords, _ = tf_inspect.getargspec(expression2)
+  if args or keywords:
+    raise ValueError('{} may not have any arguments'.format(expression2))
+
+  if tensor_util.is_tf_type(expression1):
+    return _tf_assert_stmt(expression1, expression2)
+  else:
+    return _py_assert_stmt(expression1, expression2)
+
+
+def _tf_assert_stmt(expression1, expression2):
+  """Overload of assert_stmt that stages a TF Assert.
+
+  This implementation deviates from Python semantics as follows:
+    (1) the assertion is verified regardless of the state of __debug__
+    (2) on assertion failure, the graph execution will fail with
+        tensorflow.errors.ValueError, rather than AssertionError.
+
+  Args:
+    expression1: tensorflow.Tensor, must evaluate to a tf.bool scalar
+    expression2: Callable[[], Union[tensorflow.Tensor, List[tensorflow.Tensor]]]
+
+  Returns:
+    tensorflow.Operation
+  """
+  expression2_tensors = expression2()
+  if not isinstance(expression2_tensors, list):
+    expression2_tensors = [expression2_tensors]
+  return control_flow_assert.Assert(expression1, expression2_tensors)
+
+
+def _py_assert_stmt(expression1, expression2):
+  """Overload of assert_stmt that executes a Python assert statement."""
+  assert expression1, expression2()
+  return None

tensorflow/python/autograph/operators/exceptions_test.py

@@ -0,0 +1,86 @@
+# 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 exceptions module."""
+
+from tensorflow.python.autograph.operators import exceptions
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import errors_impl
+from tensorflow.python.framework import test_util
+from tensorflow.python.platform import test
+
+
+class ExceptionsTest(test.TestCase):
+
+  def test_assert_tf_untriggered(self):
+    with self.cached_session() as sess:
+      t = exceptions.assert_stmt(
+          constant_op.constant(True), lambda: constant_op.constant('ignored'))
+      self.evaluate(t)
+
+  @test_util.run_deprecated_v1
+  def test_assert_tf_triggered(self):
+    with self.cached_session() as sess:
+      t = exceptions.assert_stmt(
+          constant_op.constant(False),
+          lambda: constant_op.constant('test message'))
+
+      with self.assertRaisesRegex(errors_impl.InvalidArgumentError,
+                                  'test message'):
+        self.evaluate(t)
+
+  @test_util.run_deprecated_v1
+  def test_assert_tf_multiple_printed_values(self):
+    two_tensors = [
+        constant_op.constant('test message'),
+        constant_op.constant('another message')
+    ]
+    with self.cached_session() as sess:
+      t = exceptions.assert_stmt(
+          constant_op.constant(False), lambda: two_tensors)
+
+      with self.assertRaisesRegex(errors_impl.InvalidArgumentError,
+                                  'test message.*another message'):
+        self.evaluate(t)
+
+  def test_assert_python_untriggered(self):
+    side_effect_trace = []
+
+    def expression_with_side_effects():
+      side_effect_trace.append(object())
+      return 'test message'
+
+    exceptions.assert_stmt(True, expression_with_side_effects)
+
+    self.assertListEqual(side_effect_trace, [])
+
+  def test_assert_python_triggered(self):
+    if not __debug__:
+      # Python assertions only be tested when in debug mode.
+      return
+
+    side_effect_trace = []
+    tracer = object()
+
+    def expression_with_side_effects():
+      side_effect_trace.append(tracer)
+      return 'test message'
+
+    with self.assertRaisesRegex(AssertionError, 'test message'):
+      exceptions.assert_stmt(False, expression_with_side_effects)
+    self.assertListEqual(side_effect_trace, [tracer])
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/autograph/operators/logical.py

@@ -0,0 +1,96 @@
+# 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.
+# ==============================================================================
+"""Logical boolean operators: not, and, or."""
+
+from tensorflow.python.framework import tensor_util
+from tensorflow.python.ops import cond as tf_cond
+from tensorflow.python.ops import gen_math_ops
+
+
+def not_(a):
+  """Functional form of "not"."""
+  if tensor_util.is_tf_type(a):
+    return _tf_not(a)
+  return _py_not(a)
+
+
+def _tf_not(a):
+  """Implementation of the "not_" operator for TensorFlow."""
+  return gen_math_ops.logical_not(a)
+
+
+def _py_not(a):
+  """Default Python implementation of the "not_" operator."""
+  return not a
+
+
+def and_(a, b):
+  """Functional form of "and". Uses lazy evaluation semantics."""
+  a_val = a()
+  if tensor_util.is_tf_type(a_val):
+    return _tf_lazy_and(a_val, b)
+  return _py_lazy_and(a_val, b)
+
+
+def _tf_lazy_and(cond, b):
+  """Lazy-eval equivalent of "and" for Tensors."""
+  # TODO(mdan): Enforce cond is scalar here?
+  return tf_cond.cond(cond, b, lambda: cond)
+
+
+def _py_lazy_and(cond, b):
+  """Lazy-eval equivalent of "and" in Python."""
+  return cond and b()
+
+
+def or_(a, b):
+  """Functional form of "or". Uses lazy evaluation semantics."""
+  a_val = a()
+  if tensor_util.is_tf_type(a_val):
+    return _tf_lazy_or(a_val, b)
+  return _py_lazy_or(a_val, b)
+
+
+def _tf_lazy_or(cond, b):
+  """Lazy-eval equivalent of "or" for Tensors."""
+  # TODO(mdan): Enforce cond is scalar here?
+  return tf_cond.cond(cond, lambda: cond, b)
+
+
+def _py_lazy_or(cond, b):
+  """Lazy-eval equivalent of "or" in Python."""
+  return cond or b()
+
+
+def eq(a, b):
+  """Functional form of "equal"."""
+  if tensor_util.is_tf_type(a) or tensor_util.is_tf_type(b):
+    return _tf_equal(a, b)
+  return _py_equal(a, b)
+
+
+def _tf_equal(a, b):
+  """Overload of "equal" for Tensors."""
+  return gen_math_ops.equal(a, b)
+
+
+def _py_equal(a, b):
+  """Overload of "equal" that falls back to Python's default implementation."""
+  return a == b
+
+
+def not_eq(a, b):
+  """Functional form of "not-equal"."""
+  return not_(eq(a, b))

tensorflow/python/autograph/operators/logical_test.py

@@ -0,0 +1,84 @@
+# 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 logical module."""
+
+from tensorflow.python.autograph.operators import logical
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import test_util
+from tensorflow.python.platform import test
+
+
+class LogicalOperatorsTest(test.TestCase):
+
+  def assertNotCalled(self):
+    self.fail('this should not be called')
+
+  def _tf_true(self):
+    return constant_op.constant(True)
+
+  def _tf_false(self):
+    return constant_op.constant(False)
+
+  def test_and_python(self):
+    self.assertTrue(logical.and_(lambda: True, lambda: True))
+    self.assertTrue(logical.and_(lambda: [1], lambda: True))
+    self.assertListEqual(logical.and_(lambda: True, lambda: [1]), [1])
+
+    self.assertFalse(logical.and_(lambda: False, lambda: True))
+    self.assertFalse(logical.and_(lambda: False, self.assertNotCalled))
+
+  @test_util.run_deprecated_v1
+  def test_and_tf(self):
+    with self.cached_session() as sess:
+      t = logical.and_(self._tf_true, self._tf_true)
+      self.assertEqual(self.evaluate(t), True)
+      t = logical.and_(self._tf_true, lambda: True)
+      self.assertEqual(self.evaluate(t), True)
+      t = logical.and_(self._tf_false, lambda: True)
+      self.assertEqual(self.evaluate(t), False)
+      # TODO(mdan): Add a test for ops with side effects.
+
+  def test_or_python(self):
+    self.assertFalse(logical.or_(lambda: False, lambda: False))
+    self.assertFalse(logical.or_(lambda: [], lambda: False))
+    self.assertListEqual(logical.or_(lambda: False, lambda: [1]), [1])
+
+    self.assertTrue(logical.or_(lambda: False, lambda: True))
+    self.assertTrue(logical.or_(lambda: True, self.assertNotCalled))
+
+  @test_util.run_deprecated_v1
+  def test_or_tf(self):
+    with self.cached_session() as sess:
+      t = logical.or_(self._tf_false, self._tf_true)
+      self.assertEqual(self.evaluate(t), True)
+      t = logical.or_(self._tf_false, lambda: True)
+      self.assertEqual(self.evaluate(t), True)
+      t = logical.or_(self._tf_true, lambda: True)
+      self.assertEqual(self.evaluate(t), True)
+      # TODO(mdan): Add a test for ops with side effects.
+
+  def test_not_python(self):
+    self.assertFalse(logical.not_(True))
+    self.assertFalse(logical.not_([1]))
+    self.assertTrue(logical.not_([]))
+
+  def test_not_tf(self):
+    with self.cached_session() as sess:
+      t = logical.not_(self._tf_false())
+      self.assertEqual(self.evaluate(t), True)
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/autograph/operators/py_builtins.py

@@ -0,0 +1,533 @@
+# 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.
+# ==============================================================================
+"""Operators corresponding to Python builtin functions.
+
+List of built-in functions: https://docs.python.org/3/library/functions.html
+"""
+
+import inspect
+
+from tensorflow.python.autograph.utils import tensors
+from tensorflow.python.autograph.utils import type_registry
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import cond
+from tensorflow.python.ops import control_flow_assert
+from tensorflow.python.ops import gen_parsing_ops
+from tensorflow.python.ops import gen_string_ops
+from tensorflow.python.ops import list_ops
+from tensorflow.python.ops import math_ops
+
+
+UNSPECIFIED = object()
+
+abs_registry = type_registry.TypeRegistry()
+len_registry = type_registry.TypeRegistry()
+print_registry = type_registry.TypeRegistry()
+enumerate_registry = type_registry.TypeRegistry()
+zip_registry = type_registry.TypeRegistry()
+map_registry = type_registry.TypeRegistry()
+filter_registry = type_registry.TypeRegistry()
+any_registry = type_registry.TypeRegistry()
+all_registry = type_registry.TypeRegistry()
+sorted_registry = type_registry.TypeRegistry()
+next_registry = type_registry.TypeRegistry()
+
+
+def registry_lookup(reg, obj):
+  try:
+    return reg.lookup(obj)
+  except LookupError:
+    pass
+  return None
+
+
+def overload_of(f):
+  if f in SUPPORTED_BUILTINS:
+    return BUILTIN_FUNCTIONS_MAP[f.__name__]
+  return f
+
+
+def _find_originating_frame(caller_fn_scope, innermost=True):
+  """Locates the frame in which `caller_fn_scope` was defined."""
+  ctx_frame = inspect.currentframe()
+  result = None
+  while ctx_frame is not None:
+    # Note it should not be normally possible to get false positives this way
+    # because the function scope object is not accessible to user code (barring
+    # call stack introspection).
+    if ctx_frame.f_locals.get(caller_fn_scope.name, None) is caller_fn_scope:
+      result = ctx_frame
+      if innermost:
+        break
+    ctx_frame = ctx_frame.f_back
+
+  assert result is not None, (
+      'the conversion process should ensure the caller_fn_scope is always'
+      ' found somewhere on the call stack')
+
+  return result
+
+
+def locals_in_original_context(caller_fn_scope):
+  """Executes the locals function in the context of a specified function."""
+  return _find_originating_frame(caller_fn_scope, innermost=True).f_locals
+
+
+def globals_in_original_context(caller_fn_scope):
+  """Executes the locals function in the context of a specified function."""
+  return _find_originating_frame(caller_fn_scope, innermost=True).f_globals
+
+
+def eval_in_original_context(f, args, caller_fn_scope):
+  """Executes the eval function in the context of a specified function."""
+  # When control flow is rewritten using functions, eval should use the
+  # variables found in the same block where it was called. That is equivalent
+  # to the innermost function call.
+  ctx_frame = _find_originating_frame(caller_fn_scope, innermost=True)
+
+  args = (
+      args[0],
+      ctx_frame.f_globals if len(args) < 2 else args[1],
+      ctx_frame.f_locals if len(args) < 3 else args[2],
+  )
+  return f(*args)
+
+
+def super_in_original_context(f, args, caller_fn_scope):
+  """Executes the super function in the context of a specified function.
+
+  See https://docs.python.org/3/library/functions.html#super for the exact
+  details
+
+  Args:
+    f: Callable, typically the super builtin
+    args: List[Any], the original call arguments
+    caller_fn_scope: Optional[function_wrappers.FunctionScope], the function
+      scope of the converted function in which this call was originally made
+
+  Returns:
+    The result of calling `f` as if it was called in the frame indicated by
+      `caller_fn_scope`.
+  """
+
+  # Only the no-arg call is desugared.
+  if args:
+    return f(*args)
+
+  # Inner functions seem to include their closure in f_locals, so we need
+  # to find the outermost frame.
+  ctx_frame = _find_originating_frame(caller_fn_scope, innermost=False)
+
+  # When super(..) is called without arguments, it looks for __class__ cell
+  # variable and the first argument passed in the enclosing function according
+  # to the spec https://www.python.org/dev/peps/pep-3135/ .
+  #
+  # We couldn't verify if `inspect.currentframe().f_code.co_varnames[0]` is
+  # guaranteed to be the first argument from an official doc or PEP, however,
+  # it's fairly stable and well established:
+  # - An unofficial community doc mentions it.
+  #   https://python-reference.readthedocs.io/en/latest/docs/code/varnames.html
+  # - CPython has tests checking that order, which was merged in 2008, and
+  #   unchanged since then.
+  #   https://github.com/python/cpython/blame/2f224a077a83ac9de8a12bb7dcc516642b8176d8/Lib/lib2to3/tests/data/py2_test_grammar.py#L157
+  #   https://github.com/python/cpython/blame/2f224a077a83ac9de8a12bb7dcc516642b8176d8/Lib/lib2to3/tests/data/py3_test_grammar.py#L192
+  #
+  # Note: the name can be more reliably obtained by inspecting the calling
+  # function's argspec.
+  #
+  # Even though methods can be declared using *args (def method(*args)),
+  # that pattern is disallowed by super() -- it raises super() no arguments.
+  # Method definitions using **kwargs are not allowed at all.
+  # In other words, we can always assume that self is on the first positional
+  # argument (for correct code).
+  #
+  # TODO(mdan): Consider additional checks in case the input code is incorrect.
+  # For example, the error might be cryptic compared to what super() regularly
+  # raises.
+
+  type_arg = ctx_frame.f_locals['__class__']
+  self_arg_name = ctx_frame.f_code.co_varnames[0]
+  self_arg = ctx_frame.f_locals[self_arg_name]
+  return f(type_arg, self_arg)
+
+
+def abs_(x):
+  abs_override = registry_lookup(abs_registry, x)
+  if abs_override is not None:
+    return abs_override(x)
+  if tensor_util.is_tf_type(x):
+    return _tf_abs(x)
+  return _py_abs(x)
+
+
+def _tf_abs(x):
+  return math_ops.abs(x)
+
+
+def _py_abs(x):
+  return abs(x)
+
+
+def float_(x=0):
+  if tensor_util.is_tf_type(x):
+    return _tf_float(x)
+  return _py_float(x)
+
+
+def _tf_float(x):
+  # TODO(mdan): We shouldn't assume float32.
+  if x.dtype == dtypes.string:
+    return gen_parsing_ops.string_to_number(x, out_type=dtypes.float32)
+  return math_ops.cast(x, dtype=dtypes.float32)
+
+
+def _py_float(x):
+  return float(x)
+
+
+def int_(x=0, base=UNSPECIFIED):
+  if tensor_util.is_tf_type(x):
+    return _tf_int(x, base)
+  return _py_int(x, base)
+
+
+def _tf_int(x, base):
+  if base not in (10, UNSPECIFIED):
+    raise NotImplementedError('base {} not supported for int'.format(base))
+
+  # TODO(mdan): We shouldn't assume int32.
+  if x.dtype == dtypes.string:
+    return gen_parsing_ops.string_to_number(x, out_type=dtypes.int32)
+  return math_ops.cast(x, dtype=dtypes.int32)
+
+
+def _py_int(x, base):
+  if base is UNSPECIFIED:
+    return int(x)
+  return int(x, base)
+
+
+def len_(s):
+  len_override = registry_lookup(len_registry, s)
+  if len_override is not None:
+    return len_override(s)
+  if tensors.is_tensor_array(s):
+    return _tf_tensor_array_len(s)
+  elif tensors.is_tensor_list(s):
+    return _tf_tensor_list_len(s)
+  elif tensor_util.is_tf_type(s):
+    return _tf_tensor_len(s)
+  return _py_len(s)
+
+
+def _tf_tensor_array_len(s):
+  return s.size()
+
+
+def _tf_tensor_list_len(s):
+  return list_ops.tensor_list_length(s)
+
+
+def _tf_tensor_len(s):
+  """Overload of len_ for Tensor arguments."""
+  # Statically shaped tensors: length is known ahead of time.
+  if s.shape.ndims and s.shape.dims[0].value is not None:
+    return s.shape.dims[0].value
+
+  # Static shape of unknown dimensions: use dynamic shape but statically
+  # check that it's a scalar.
+  shape = array_ops.shape(s)
+
+  assert shape.shape, 'shape tensor of zero size? {}'.format(shape)
+
+  if shape.shape[0] == 0:
+    raise ValueError(
+        'len requires a non-scalar tensor, got one of shape {}'.format(shape))
+
+  if shape.shape.dims[0].value is not None:
+    return array_ops.shape(s)[0]
+
+  # Fully dynamic shape: use ops.
+  rank = array_ops.rank(s)
+
+  def raise_zero_rank_error():
+    msg = gen_string_ops.string_join(
+        ['len requires non-zero rank, got ',
+         gen_string_ops.as_string(rank)])
+    with ops.control_dependencies([control_flow_assert.Assert(False, [msg])]):
+      return constant_op.constant(0, dtype=dtypes.int32)
+
+  return cond.cond(rank > 0, lambda: array_ops.shape(s)[0],
+                   raise_zero_rank_error)
+
+
+def _py_len(s):
+  return len(s)
+
+
+def print_(*objects, **kwargs):
+  """Overload of the print builtin."""
+  # Note: Python 2.6 doesn't support explicit keywords after starargs.
+  unknown_kwargs = tuple(
+      set(kwargs.keys()) - set(('sep', 'end', 'file', 'flush')))
+  if unknown_kwargs:
+    raise ValueError('invalid keyword arguments: {}'.format(unknown_kwargs))
+
+  print_fn = _py_print
+  for x in objects:
+    print_override = registry_lookup(print_registry, x)
+    if print_override is not None:  # pylint: disable=comparison-with-callable
+      print_fn = print_override
+      break
+
+  if print_fn is _py_print:
+    # If this fails, ops/autograph_ops.py hasn't been imported.
+    assert not any(tensor_util.is_tf_type(s) for s in objects)
+
+  return print_fn(*objects, **kwargs)
+
+
+def _py_print(*objects, **kwargs):
+  print(*objects, **kwargs)
+
+
+def min_(*args, **kwargs):
+  if any(tensor_util.is_tf_type(s) for s in args):
+    return _tf_min(*args, **kwargs)
+  return _py_min(*args, **kwargs)
+
+
+def _tf_min(*args, **kwargs):
+  if len(kwargs):
+    kwargs_tuple = tuple(set(kwargs.keys()))
+    raise ValueError('These keyword arguments are '
+                     'currently not supported: {}'.format(kwargs_tuple))
+  if len(args) == 1:
+    rank = args[0].shape.rank
+    if rank == 0:
+      return args[0]
+    if rank == 1:
+      return math_ops.reduce_min(*args, axis=0)
+    raise ValueError('min(arg) currently support only tensor with rank 1, '
+                     'but got a tensor with rank {}'.format(rank))
+  for arg in args:
+    rank = arg.shape.rank
+    if rank != 0:
+      raise ValueError('min(arg1, arg2, *args) currently support '
+                       'only scalar tensor, but got a tensor '
+                       'with shape {}'.format(rank))
+  return math_ops.reduce_min(args, axis=0)
+
+
+def _py_min(*args, **kwargs):
+  return min(*args, **kwargs)
+
+
+def max_(*args, **kwargs):
+  if any(tensor_util.is_tf_type(s) for s in args):
+    return _tf_max(*args, **kwargs)
+  return _py_max(*args, **kwargs)
+
+
+def _tf_max(*args, **kwargs):
+  if len(kwargs):
+    kwargs_tuple = tuple(set(kwargs.keys()))
+    raise ValueError('These keyword arguments are '
+                     'currently not supported: {}'.format(kwargs_tuple))
+  if len(args) == 1:
+    rank = args[0].shape.rank
+    if rank == 0:
+      return args[0]
+    if rank == 1:
+      return math_ops.reduce_max(*args, axis=0)
+    raise ValueError('max(arg) currently support only tensor with rank 1, '
+                     'but got a tensor with rank {}'.format(rank))
+  for arg in args:
+    rank = arg.shape.rank
+    if rank != 0:
+      raise ValueError('max(arg1, arg2, *args) currently support '
+                       'only scalar tensor, but got a tensor '
+                       'with shape {}'.format(rank))
+  return math_ops.reduce_max(args, axis=0)
+
+
+def _py_max(*args, **kwargs):
+  return max(*args, **kwargs)
+
+
+def range_(start_or_stop, stop=UNSPECIFIED, step=UNSPECIFIED):
+  if any(tensor_util.is_tf_type(s) for s in (start_or_stop, stop, step)):
+    return _tf_range(start_or_stop, stop, step)
+  return _py_range(start_or_stop, stop, step)
+
+
+def _tf_range(start_or_stop, stop, step):
+  """Overload of range_ that generates a TF range tensor."""
+  # Note: for static inputs (e.g. constants), tf.range errors out at graph
+  # construction time, instead of returning an empty tensor. Preventing the
+  # graph construction error aligns the semantics with Python.
+
+  # TODO(mdan): We should optimize this when a full tensor is not required.
+  if step is not UNSPECIFIED:
+    # TODO(mdan): Add argument coercion similar to other cases.
+    return math_ops.range(start_or_stop, stop, step)
+  if stop is not UNSPECIFIED:
+    stop = math_ops.maximum(start_or_stop, stop)
+    return math_ops.range(start_or_stop, stop)
+  start_or_stop = math_ops.maximum(start_or_stop, 0)
+  return math_ops.range(start_or_stop)
+
+
+def _py_range(start_or_stop, stop, step):
+  if step is not UNSPECIFIED:
+    return range(start_or_stop, stop, step)
+  if stop is not UNSPECIFIED:
+    return range(start_or_stop, stop)
+  return range(start_or_stop)
+
+
+def enumerate_(s, start=0):
+  enumerate_override = registry_lookup(enumerate_registry, s)
+  if enumerate_override is not None:
+    return enumerate_override(s, start)
+  return _py_enumerate(s, start)
+
+
+def _py_enumerate(s, start=0):
+  return enumerate(s, start)
+
+
+def zip_(*iterables, strict=False):
+  zip_fn = _py_zip
+  # If the overridden function is not the same across all iterables, use _py_zip
+  for x in iterables:
+    zip_override = registry_lookup(zip_registry, x)
+    if zip_override is None or (zip_fn != _py_zip and zip_override != zip_fn):  # pylint: disable=comparison-with-callable
+      zip_fn = _py_zip
+      break
+    zip_fn = zip_override
+  return zip_fn(*iterables, strict=strict)
+
+
+def _py_zip(*iterables, strict=False):
+  if strict:
+    return zip(*iterables, strict=True)
+  else:
+    # Python < 3.10 doesn't have `strict` kwarg.
+    return zip(*iterables)
+
+
+def map_(fn, *iterables):
+  map_fn = _py_map
+  # If the overridden function is not the same across all iterables, use _py_map
+  for x in iterables:
+    map_override = registry_lookup(map_registry, x)
+    if map_override is None or (map_fn != _py_map and map_override != map_fn):  # pylint: disable=comparison-with-callable
+      map_fn = _py_map
+      break
+    map_fn = map_override
+  return map_fn(fn, *iterables)
+
+
+def _py_map(fn, *iterables):
+  return map(fn, *iterables)
+
+
+def next_(iterator, default=UNSPECIFIED):
+  next_override = registry_lookup(next_registry, iterator)
+  if next_override is not None:
+    return next_override(iterator, default)
+  return next_py(iterator, default)
+
+
+def next_py(iterator, default=UNSPECIFIED):
+  if default is UNSPECIFIED:
+    return next(iterator)
+  return next(iterator, default)
+
+
+def filter_(function, iterable):
+  filter_override = registry_lookup(filter_registry, iterable)
+  if filter_override is not None:
+    return filter_override(function, iterable)
+  return _py_filter(function, iterable)
+
+
+def _py_filter(function, iterable):
+  return filter(function, iterable)
+
+
+def any_(iterable):
+  any_override = registry_lookup(any_registry, iterable)
+  if any_override is not None:
+    return any_override(iterable)
+  return _py_any(iterable)
+
+
+def _py_any(iterable):
+  return any(iterable)
+
+
+def all_(iterable):
+  all_override = registry_lookup(all_registry, iterable)
+  if all_override is not None:
+    return all_override(iterable)
+  return _py_all(iterable)
+
+
+def _py_all(iterable):
+  return all(iterable)
+
+
+def sorted_(iterable, key=UNSPECIFIED, reverse=UNSPECIFIED):
+  sorted_override = registry_lookup(sorted_registry, iterable)
+  if sorted_override is not None:
+    return sorted_override(iterable, key, reverse)
+  return _py_sorted(iterable, key, reverse)
+
+
+def _py_sorted(iterable, key, reverse):
+  if key is not UNSPECIFIED and reverse is UNSPECIFIED:
+    return sorted(iterable, key=key)
+  if key is UNSPECIFIED and reverse is not UNSPECIFIED:
+    return sorted(iterable, reverse=reverse)
+  if key is not UNSPECIFIED and reverse is not UNSPECIFIED:
+    return sorted(iterable, key=key, reverse=reverse)
+  return sorted(iterable)
+
+
+SUPPORTED_BUILTINS = (abs, float, int, len, print, range, enumerate, zip, map,
+                      filter, any, all, sorted)
+
+BUILTIN_FUNCTIONS_MAP = {
+    'abs': abs_,
+    'any': any_,
+    'all': all_,
+    'enumerate': enumerate_,
+    'filter': filter_,
+    'float': float_,
+    'int': int_,
+    'len': len_,
+    'map': map_,
+    'next': next_,
+    'print': print_,
+    'range': range_,
+    'sorted': sorted_,
+    'zip': zip_,
+}

tensorflow/python/autograph/operators/py_builtins_test.py

@@ -0,0 +1,718 @@
+# 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 py_builtins module."""
+
+import io
+import sys
+
+
+from tensorflow.python.autograph.core import converter
+from tensorflow.python.autograph.core import function_wrappers
+from tensorflow.python.autograph.operators import data_structures
+from tensorflow.python.autograph.operators import py_builtins
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors_impl
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import tensor_array_ops
+from tensorflow.python.platform import test
+
+
+class TestBase:
+
+  def overridden_method(self, x):
+    return x + 20
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class PyBuiltinsTest(test.TestCase):
+
+  def test_abs(self):
+    self.assertEqual(py_builtins.abs_(-1), 1)
+    with self.cached_session() as sess:
+      t = py_builtins.abs_(constant_op.constant(-1))
+      self.assertEqual(self.evaluate(t), 1)
+      t = py_builtins.abs_(constant_op.constant([-1, 2, -3]))
+      self.assertAllEqual(self.evaluate(t), [1, 2, 3])
+
+  def test_abs_dataset(self):
+    dataset = dataset_ops.DatasetV2.from_tensor_slices([-1, 2, 3])
+    dataset = py_builtins.abs_(dataset)
+    iterator = dataset_ops.make_one_shot_iterator(dataset)
+    with self.cached_session() as sess:
+      self.assertAllEqual(self.evaluate(iterator.get_next()), 1)
+      self.assertAllEqual(self.evaluate(iterator.get_next()), 2)
+      self.assertAllEqual(self.evaluate(iterator.get_next()), 3)
+
+  def test_abs_dataset_zipped(self):
+    dataset_1 = dataset_ops.DatasetV2.from_tensor_slices([-1, 2, 3])
+    dataset_2 = dataset_ops.DatasetV2.from_tensor_slices([1, -2, 3])
+    dataset = dataset_ops.DatasetV2.zip((dataset_1, dataset_2))
+    dataset = py_builtins.abs_(dataset)
+    iterator = dataset_ops.make_one_shot_iterator(dataset)
+    with self.cached_session() as sess:
+      self.assertAllEqual(self.evaluate(iterator.get_next()), (1, 1))
+      self.assertAllEqual(self.evaluate(iterator.get_next()), (2, 2))
+      self.assertAllEqual(self.evaluate(iterator.get_next()), (3, 3))
+
+  def test_abs_dataset_mixed(self):
+    dataset_1 = dataset_ops.DatasetV2.from_tensor_slices([-1, 2, 3])
+    dataset_2 = dataset_ops.DatasetV2.from_tensor_slices([1, -2, 3])
+    dataset_3 = dataset_ops.DatasetV2.from_tensor_slices([-1, -2, -3])
+    dataset_4 = dataset_ops.DatasetV2.zip((dataset_1, dataset_2))
+    dataset = dataset_ops.DatasetV2.zip((dataset_3, dataset_4))
+    dataset = py_builtins.abs_(dataset)
+    iterator = dataset_ops.make_one_shot_iterator(dataset)
+    with self.cached_session() as sess:
+      for i in range(1, 4):
+        actual = self.evaluate(iterator.get_next())
+        self.assertAllEqual(actual[0], i)
+        self.assertAllEqual(actual[1], (i, i))
+
+  def test_float(self):
+    self.assertEqual(py_builtins.float_(10), 10.0)
+    self.assertEqual(py_builtins.float_('10.0'), 10.0)
+    with self.cached_session() as sess:
+      t = py_builtins.float_(constant_op.constant(1, dtype=dtypes.int64))
+      self.assertEqual(self.evaluate(t), 1.0)
+      st = py_builtins.float_(constant_op.constant('1.0'))
+      self.assertEqual(self.evaluate(st), 1.0)
+
+  def test_int(self):
+    self.assertEqual(py_builtins.int_(10.0), 10)
+    self.assertEqual(py_builtins.int_('11', 2), 3)
+    with self.cached_session() as sess:
+      t = py_builtins.int_(constant_op.constant(1, dtype=dtypes.float64))
+      self.assertEqual(self.evaluate(t), 1)
+      st = py_builtins.int_(constant_op.constant('1'))
+      self.assertEqual(self.evaluate(st), 1)
+      st = py_builtins.int_(constant_op.constant('1'), 10)
+      self.assertEqual(self.evaluate(st), 1)
+
+  def test_int_unsupported_base(self):
+    t = constant_op.constant(1, dtype=dtypes.float64)
+    with self.assertRaises(NotImplementedError):
+      py_builtins.int_(t, 2)
+
+  def test_len(self):
+    self.assertEqual(py_builtins.len_([1, 2, 3]), 3)
+    with self.cached_session() as sess:
+      t = py_builtins.len_(constant_op.constant([[1], [2], [3]]))
+      self.assertEqual(t, 3)
+      ta = py_builtins.len_(tensor_array_ops.TensorArray(dtypes.int32, size=5))
+      self.assertEqual(self.evaluate(ta), 5)
+      tl = py_builtins.len_(data_structures.tf_tensor_list_new([3, 4, 5]))
+      self.assertEqual(self.evaluate(tl), 3)
+
+  def test_len_dataset(self):
+    dataset = dataset_ops.DatasetV2.from_tensor_slices([3, 2, 1])
+    self.assertEqual(self.evaluate(py_builtins.len_(dataset)), 3)
+
+    # graph mode
+    @def_function.function(autograph=False)
+    def test_fn():
+      dataset = dataset_ops.DatasetV2.from_tensor_slices([3, 2, 1])
+      return py_builtins.len_(dataset)
+
+    self.assertEqual(self.evaluate(test_fn()), 3)
+
+  def test_len_dataset_infinite(self):
+    dataset = dataset_ops.DatasetV2.range(5).repeat().batch(2)
+    with self.assertRaises(errors_impl.InvalidArgumentError):
+      _ = self.evaluate(py_builtins.len_(dataset))
+
+    # graph mode
+    @def_function.function
+    def test_fn():
+      dataset = dataset_ops.DatasetV2.range(5).repeat().batch(2)
+      return py_builtins.len_(dataset)
+
+    with self.assertRaises(errors_impl.InvalidArgumentError):
+      self.evaluate(test_fn())
+
+  def test_len_dataset_unknown(self):
+    dataset = dataset_ops.DatasetV2.range(5).filter(lambda _: True).batch(2)
+    with self.assertRaises(errors_impl.InvalidArgumentError):
+      _ = self.evaluate(py_builtins.len_(dataset))
+
+    # graph mode
+    @def_function.function(autograph=False)
+    def test_fn():
+      dataset = dataset_ops.DatasetV2.range(5).filter(lambda _: True).batch(2)
+      return py_builtins.len_(dataset)
+
+    with self.assertRaises(errors_impl.InvalidArgumentError):
+      self.evaluate(test_fn())
+
+  def test_len_scalar(self):
+    with self.assertRaises(ValueError):
+      py_builtins.len_(constant_op.constant(1))
+
+  @test_util.run_deprecated_v1
+  def test_len_dynamic_shape(self):
+    with self.cached_session() as sess:
+      p = array_ops.placeholder(dtype=dtypes.int32, shape=None)
+      t = py_builtins.len_(p)
+      self.assertEqual(sess.run(t, {p: [1, 2, 3]}), 3)
+
+      with self.assertRaises(errors_impl.InvalidArgumentError):
+        t = py_builtins.len_(p)
+        sess.run(t, {p: 1})
+
+  @test_util.run_deprecated_v1
+  def test_print_tensors(self):
+    try:
+      out_capturer = io.StringIO()
+      sys.stdout = out_capturer
+      with self.cached_session() as sess:
+        sess.run(py_builtins.print_(constant_op.constant('test message'), 1))
+        self.assertEqual(out_capturer.getvalue(), 'test message 1\n')
+    finally:
+      sys.stdout = sys.__stdout__
+
+  @test_util.run_deprecated_v1
+  def test_print_complex(self):
+    try:
+      out_capturer = io.StringIO()
+      sys.stdout = out_capturer
+      with self.cached_session() as sess:
+        sess.run(
+            py_builtins.print_(constant_op.constant('test message'), [1, 2]))
+        self.assertEqual(out_capturer.getvalue(), 'test message [1, 2]\n')
+    finally:
+      sys.stdout = sys.__stdout__
+
+  def test_max(self):
+    self.assertEqual(py_builtins.max_([1, 3, 2]), 3)
+    self.assertEqual(py_builtins.max_(0, 2, 1), 2)
+
+  def test_max_tensor(self):
+    r = py_builtins.max_(constant_op.constant(2))
+    self.assertAllEqual(self.evaluate(r), 2)
+    with self.assertRaises(ValueError):
+      py_builtins.max_(constant_op.constant([[2]]))
+    r = py_builtins.max_(constant_op.constant([1, 3, 2]))
+    self.assertAllEqual(self.evaluate(r), 3)
+    with self.assertRaises(ValueError):
+      py_builtins.max_(constant_op.constant([[1, 3], [3, 4]]))
+    r = py_builtins.max_(
+        constant_op.constant(6), constant_op.constant(4),
+        constant_op.constant(8))
+    self.assertAllEqual(self.evaluate(r), 8)
+    with self.assertRaises(ValueError):
+      py_builtins.max_(
+          constant_op.constant([6]), constant_op.constant(4),
+          constant_op.constant(8))
+
+  def test_min(self):
+    self.assertEqual(py_builtins.min_([2, 1, 3]), 1)
+    self.assertEqual(py_builtins.min_(2, 0, 1), 0)
+
+  def test_min_tensor(self):
+    r = py_builtins.min_(constant_op.constant(2))
+    self.assertAllEqual(self.evaluate(r), 2)
+    with self.assertRaises(ValueError):
+      py_builtins.min_(constant_op.constant([[2]]))
+    r = py_builtins.min_(constant_op.constant([3, 1, 2]))
+    self.assertAllEqual(self.evaluate(r), 1)
+    with self.assertRaises(ValueError):
+      py_builtins.min_(constant_op.constant([[1, 3], [3, 4]]))
+    r = py_builtins.min_(
+        constant_op.constant(6), constant_op.constant(4),
+        constant_op.constant(8))
+    self.assertAllEqual(self.evaluate(r), 4)
+    with self.assertRaises(ValueError):
+      py_builtins.min_(
+          constant_op.constant([6]), constant_op.constant(4),
+          constant_op.constant(8))
+
+
+  def test_range(self):
+    self.assertListEqual(list(py_builtins.range_(3)), [0, 1, 2])
+    self.assertListEqual(list(py_builtins.range_(1, 3)), [1, 2])
+    self.assertListEqual(list(py_builtins.range_(2, 0, -1)), [2, 1])
+
+  def test_range_tensor(self):
+    with self.cached_session() as sess:
+      r = py_builtins.range_(constant_op.constant(3))
+      self.assertAllEqual(self.evaluate(r), [0, 1, 2])
+      r = py_builtins.range_(1, constant_op.constant(3))
+      self.assertAllEqual(self.evaluate(r), [1, 2])
+      r = py_builtins.range_(2, 0, constant_op.constant(-1))
+      self.assertAllEqual(self.evaluate(r), [2, 1])
+
+  def test_range_tensor_empty_range(self):
+    with self.session() as sess:
+      r = py_builtins.range_(constant_op.constant(-3))
+      self.assertAllEqual(self.evaluate(r), [])
+      r = py_builtins.range_(5, constant_op.constant(2))
+      self.assertAllEqual(self.evaluate(r), [])
+
+  def test_enumerate(self):
+    self.assertListEqual(
+        list(py_builtins.enumerate_([3, 2, 1])), [(0, 3), (1, 2), (2, 1)])
+    self.assertListEqual(
+        list(py_builtins.enumerate_([3, 2, 1], 5)), [(5, 3), (6, 2), (7, 1)])
+    self.assertListEqual(list(py_builtins.enumerate_([-8], -3)), [(-3, -8)])
+
+  def test_enumerate_dataset(self):
+    dataset = dataset_ops.DatasetV2.from_tensor_slices(['a', 'c'])
+    start = constant_op.constant(20, dtype=dtypes.int64)
+    dataset = py_builtins.enumerate_(dataset, start)
+    iterator = dataset_ops.make_one_shot_iterator(dataset)
+    with self.cached_session() as sess:
+      self.assertAllEqual(self.evaluate(iterator.get_next()), (20, b'a'))
+      self.assertAllEqual(self.evaluate(iterator.get_next()), (21, b'c'))
+
+  def test_zip(self):
+    self.assertListEqual(
+        list(py_builtins.zip_([3, 2, 1], [1, 2, 3])), [(3, 1), (2, 2), (1, 3)])
+    self.assertListEqual(
+        list(py_builtins.zip_([4, 5, 6], [-1, -2])), [(4, -1), (5, -2)])
+
+  def test_zip_dataset(self):
+    ds1 = dataset_ops.DatasetV2.from_tensor_slices([-11, -12, 4])
+    ds2 = dataset_ops.DatasetV2.from_tensor_slices([-21, -22, 5])
+    ds3 = py_builtins.zip_(ds1, ds2)
+    iterator = dataset_ops.make_one_shot_iterator(ds3)
+    with self.cached_session() as sess:
+      self.assertAllEqual(self.evaluate(iterator.get_next()), (-11, -21))
+      self.assertAllEqual(self.evaluate(iterator.get_next()), (-12, -22))
+      self.assertAllEqual(self.evaluate(iterator.get_next()), (4, 5))
+
+  def test_map(self):
+
+    def increment(x):
+      return x + 1
+
+    add_list = lambda x, y: x + y
+    self.assertListEqual(
+        list(py_builtins.map_(increment, [4, 5, 6])), [5, 6, 7])
+    self.assertListEqual(
+        list(py_builtins.map_(add_list, [3, 2, 1], [-1, -2, -3])), [2, 0, -2])
+
+  def test_map_dataset(self):
+
+    def increment(x):
+      return x + 1
+
+    ds1 = dataset_ops.DatasetV2.from_tensor_slices([4, 5, 6])
+    ds2 = py_builtins.map_(increment, ds1)
+    iterator = dataset_ops.make_one_shot_iterator(ds2)
+    with self.cached_session() as sess:
+      self.assertAllEqual(self.evaluate(iterator.get_next()), 5)
+      self.assertAllEqual(self.evaluate(iterator.get_next()), 6)
+      self.assertAllEqual(self.evaluate(iterator.get_next()), 7)
+
+  def test_map_multiple_datasets(self):
+    add_list = lambda x, y: x + y
+    ds1 = dataset_ops.DatasetV2.from_tensor_slices([-11, -12, 4])
+    ds2 = dataset_ops.DatasetV2.from_tensor_slices([-21, -22, 5])
+    ds3 = py_builtins.map_(add_list, ds1, ds2)
+    iterator = dataset_ops.make_one_shot_iterator(ds3)
+    with self.cached_session() as sess:
+      self.assertAllEqual(self.evaluate(iterator.get_next()), -32)
+      self.assertAllEqual(self.evaluate(iterator.get_next()), -34)
+      self.assertAllEqual(self.evaluate(iterator.get_next()), 9)
+
+  def test_next_normal(self):
+    iterator = iter([1, 2, 3])
+    self.assertEqual(py_builtins.next_(iterator), 1)
+    self.assertEqual(py_builtins.next_(iterator), 2)
+    self.assertEqual(py_builtins.next_(iterator), 3)
+    with self.assertRaises(StopIteration):
+      py_builtins.next_(iterator)
+    self.assertEqual(py_builtins.next_(iterator, 4), 4)
+
+  def test_next_tf_iterator(self):
+    # graph-mode iterators are only supported inside tf.function.
+    @def_function.function(autograph=False)
+    def test_fn(go_out_of_range, with_default):
+      iterator = iter(dataset_ops.Dataset.range(3))
+      retval = (
+          py_builtins.next_(iterator),
+          py_builtins.next_(iterator),
+          py_builtins.next_(iterator),
+      )
+      if go_out_of_range:
+        if with_default:
+          retval += (
+              py_builtins.next_(iterator,
+                                constant_op.constant(-3, dtype=dtypes.int64)),
+              py_builtins.next_(iterator,
+                                constant_op.constant(-4, dtype=dtypes.int64)),
+          )
+        else:
+          py_builtins.next_(iterator)
+      return retval
+
+    self.assertAllEqual(
+        self.evaluate(test_fn(go_out_of_range=False, with_default=None)),
+        (0, 1, 2))
+    self.assertAllEqual(
+        self.evaluate(test_fn(go_out_of_range=True, with_default=True)),
+        (0, 1, 2, -3, -4))
+    with self.assertRaises(errors_impl.OutOfRangeError):
+      self.evaluate(test_fn(go_out_of_range=True, with_default=False))
+
+  def test_next_tf_iterator_error_checking(self):
+    # graph-mode iterators are only supported inside tf.function.
+    @def_function.function(autograph=False)
+    def test_fn():
+      iterator = iter(dataset_ops.Dataset.range(1))
+      py_builtins.next_(iterator)
+      py_builtins.next_(iterator, constant_op.constant(-3))
+
+    # Dataset.range defaults to int64,
+    with self.assertRaisesRegex(TypeError, 'default.*int64'):
+      self.evaluate(test_fn())
+
+  def test_next_tf_iterator_error_checking_structures(self):
+    # graph-mode iterators are only supported inside tf.function.
+    @def_function.function(autograph=False)
+    def test_fn(default_val):
+      ds = dataset_ops.Dataset.range(1)
+      ds = ds.map(lambda i: {'a': i + 1, 'b': i + 10})
+      iterator = iter(ds)
+      py_builtins.next_(iterator)
+      py_builtins.next_(iterator, default_val)
+
+    default = {
+        'a': constant_op.constant(3, dtype=dtypes.int64),
+    }
+    with self.assertRaisesRegex(TypeError, 'same element structure'):
+      test_fn(default)
+    default = {
+        'a': constant_op.constant(3.0),
+        'b': [constant_op.constant(30), constant_op.constant(300)]
+    }
+    with self.assertRaisesRegex(TypeError, 'same element structure'):
+      test_fn(default)
+    default = {
+        'a': constant_op.constant(3.0),
+        'b': constant_op.constant(30, dtype=dtypes.int64),
+    }
+    with self.assertRaisesRegex(TypeError, 'float32'):
+      test_fn(default)
+
+  def _basic_function_scope(self):
+    return function_wrappers.FunctionScope(
+        'test_function_name',
+        'test_scope',  # Note: this must match the name in the `with` statement.
+        converter.ConversionOptions())
+
+  def test_eval_in_original_context(self):
+
+    def test_fn():
+      l = 1  # pylint:disable=unused-variable
+      with self._basic_function_scope() as test_scope:
+        return py_builtins.eval_in_original_context(eval, ('l',), test_scope)
+
+    self.assertEqual(test_fn(), 1)
+
+  def test_eval_in_original_context_inner_function(self):
+
+    def test_fn():
+      l = 1  # pylint:disable=unused-variable
+      with self._basic_function_scope() as test_scope:
+
+        def inner_fn():
+          # Note: a user function without a top-level function scope should
+          # never be found in user code; it's only possible in generated code.
+          l = 2  # pylint:disable=unused-variable
+          return py_builtins.eval_in_original_context(eval, ('l',), test_scope)
+
+        return inner_fn()
+
+    self.assertEqual(test_fn(), 2)
+
+  def test_locals_in_original_context(self):
+
+    def test_fn():
+      l = 1  # pylint:disable=unused-variable
+      with self._basic_function_scope() as test_scope:
+        return py_builtins.locals_in_original_context(test_scope)
+
+    locs = test_fn()
+
+    self.assertEqual(locs['l'], 1)
+
+  def test_locals_in_original_context_inner_function(self):
+
+    def test_fn():
+      l = 1  # pylint:disable=unused-variable
+      with self._basic_function_scope() as test_scope:
+
+        def inner_fn():
+          # Note: a user function without a top-level function scope should
+          # never be found in user code; it's only possible in generated code.
+          l = 2  # pylint:disable=unused-variable
+          return py_builtins.locals_in_original_context(test_scope)
+
+        return inner_fn()
+
+    locs = test_fn()
+
+    self.assertEqual(locs['l'], 2)
+
+  def test_globals_in_original_context(self):
+
+    def test_fn():
+      with self._basic_function_scope() as test_scope:
+        return py_builtins.globals_in_original_context(test_scope)
+
+    globs = test_fn()
+
+    self.assertIs(globs['TestBase'], TestBase)
+
+  def test_globals_in_original_context_inner_function(self):
+
+    def test_fn():
+      with self._basic_function_scope() as test_scope:
+
+        def inner_fn():
+          # Note: a user function without a top-level function scope should
+          # never be found in user code; it's only possible in generated code.
+          return py_builtins.globals_in_original_context(test_scope)
+
+        return inner_fn()
+
+    globs = test_fn()
+
+    self.assertIs(globs['TestBase'], TestBase)
+
+  def test_super_in_original_context_unary_call(self):
+    test_case_self = self
+
+    class TestSubclass(TestBase):
+
+      def overridden_method(self, x):
+        test_case_self.fail('This should never be called.')
+
+      def test_method(self):
+        with test_case_self._basic_function_scope() as test_scope:
+          test_base_unbound = py_builtins.super_in_original_context(
+              super, (TestSubclass,), test_scope)
+          test_base = test_base_unbound.__get__(self, TestSubclass)
+          return test_base.overridden_method(1)
+
+    tc = TestSubclass()
+    self.assertEqual(tc.test_method(), 21)
+
+  def test_super_in_original_context_binary_call(self):
+    test_case_self = self
+
+    class TestSubclass(TestBase):
+
+      def overridden_method(self, x):
+        test_case_self.fail('This should never be called.')
+
+      def test_method(self):
+        with test_case_self._basic_function_scope() as test_scope:
+          test_base = py_builtins.super_in_original_context(
+              super, (TestSubclass, self), test_scope)
+          return test_base.overridden_method(1)
+
+    tc = TestSubclass()
+    self.assertEqual(tc.test_method(), 21)
+
+  def test_super_in_original_context_niladic_call(self):
+    test_case_self = self
+
+    class TestSubclass(TestBase):
+
+      def overridden_method(self, x):
+        test_case_self.fail('This should never be called.')
+
+      def test_method(self):
+        with test_case_self._basic_function_scope() as test_scope:
+          b = py_builtins.super_in_original_context(super, (), test_scope)
+          return b.overridden_method(1)
+
+    tc = TestSubclass()
+    self.assertEqual(tc.test_method(), 21)
+
+  def test_super_in_original_context_caller_with_locals(self):
+    test_case_self = self
+
+    class TestSubclass(TestBase):
+
+      def overridden_method(self, x):
+        test_case_self.fail('This should never be called.')
+
+      def test_method(self, x):
+        y = 7
+        with test_case_self._basic_function_scope() as test_scope:
+          z = 7
+          return py_builtins.super_in_original_context(
+              super, (), test_scope).overridden_method(x + y - z)
+
+    tc = TestSubclass()
+    self.assertEqual(tc.test_method(1), 21)
+
+  def test_super_in_original_context_inner_function(self):
+    test_case_self = self
+
+    class TestSubclass(TestBase):
+
+      def overridden_method(self, x):
+        test_case_self.fail('This should never be called.')
+
+      def test_method(self, x):
+        with test_case_self._basic_function_scope() as test_scope:
+          # Oddly, it's sufficient to use `self` in an inner function
+          # to gain access to __class__ in this scope.
+          # TODO(mdan): Is this true across implementations?
+          # Note: normally, it's illegal to use super() in inner functions (it
+          # throws an error), but the generated code may create them.
+          def inner_fn():
+            return py_builtins.super_in_original_context(
+                super, (), test_scope).overridden_method(x)
+
+          return inner_fn()
+
+    tc = TestSubclass()
+    self.assertEqual(tc.test_method(1), 21)
+
+  def test_super_in_original_context_inner_lambda(self):
+    test_case_self = self
+
+    class TestSubclass(TestBase):
+
+      def overridden_method(self, x):
+        test_case_self.fail('This should never be called.')
+
+      def test_method(self, x):
+        with test_case_self._basic_function_scope() as test_scope:
+          # Oddly, it's sufficient to use `self` in an inner function
+          # to gain access to __class__ in this scope.
+          # TODO(mdan): Is this true across implementations?
+          # Note: normally, it's illegal to use super() in inner functions (it
+          # throws an error), but the generated code may create them.
+          l = lambda: py_builtins.super_in_original_context(  # pylint:disable=g-long-lambda
+              super, (), test_scope).overridden_method(x)
+          return l()
+
+    tc = TestSubclass()
+    self.assertEqual(tc.test_method(1), 21)
+
+  def test_filter(self):
+    self.assertListEqual(
+        list(py_builtins.filter_(lambda x: x == 'b', ['a', 'b', 'c'])), ['b'])
+    self.assertListEqual(
+        list(py_builtins.filter_(lambda x: x < 3, [3, 2, 1])), [2, 1])
+
+  def test_filter_dataset(self):
+    dataset = dataset_ops.DatasetV2.from_tensor_slices([3, 2, 1])
+    dataset = py_builtins.filter_(lambda x: x < 3, dataset)
+    iterator = dataset_ops.make_one_shot_iterator(dataset)
+    with self.cached_session() as sess:
+      self.assertAllEqual(self.evaluate(iterator.get_next()), 2)
+      self.assertAllEqual(self.evaluate(iterator.get_next()), 1)
+
+  def test_any(self):
+    self.assertEqual(py_builtins.any_([False, True, False]), True)
+    self.assertEqual(py_builtins.any_([False, False, False]), False)
+
+  def test_any_dataset(self):
+    dataset_1 = dataset_ops.DatasetV2.from_tensor_slices([False, True, False])
+    dataset_2 = dataset_ops.DatasetV2.from_tensor_slices([False, False, False])
+    self.assertEqual(self.evaluate(py_builtins.any_(dataset_1)), True)
+    self.assertEqual(self.evaluate(py_builtins.any_(dataset_2)), False)
+
+    dataset_3 = dataset_ops.DatasetV2.from_tensor_slices([0, 1, 2])
+    with self.assertRaises(ValueError):
+      py_builtins.any_(dataset_3)
+
+    dataset_4 = dataset_ops.DatasetV2.from_tensor_slices([False, True, False])
+    dataset_zipped = dataset_ops.DatasetV2.zip((dataset_4, dataset_4))
+    with self.assertRaises(ValueError):
+      py_builtins.any_(dataset_zipped)
+
+    dataset_mixed = dataset_ops.DatasetV2.zip((dataset_3, dataset_4))
+    with self.assertRaises(ValueError):
+      py_builtins.any_(dataset_mixed)
+
+  def test_all(self):
+    self.assertEqual(py_builtins.all_([False, True, False]), False)
+    self.assertEqual(py_builtins.all_([True, True, True]), True)
+
+  def test_all_dataset(self):
+    dataset_1 = dataset_ops.DatasetV2.from_tensor_slices([False, True, False])
+    dataset_2 = dataset_ops.DatasetV2.from_tensor_slices([True, True, True])
+    self.assertEqual(self.evaluate(py_builtins.all_(dataset_1)), False)
+    self.assertEqual(self.evaluate(py_builtins.all_(dataset_2)), True)
+
+    dataset_3 = dataset_ops.DatasetV2.from_tensor_slices([0, 1, 2])
+    with self.assertRaises(ValueError):
+      py_builtins.all_(dataset_3)
+
+    dataset_4 = dataset_ops.DatasetV2.from_tensor_slices([False, True, False])
+    dataset_zipped = dataset_ops.DatasetV2.zip((dataset_4, dataset_4))
+    with self.assertRaises(ValueError):
+      py_builtins.all_(dataset_zipped)
+
+    dataset_mixed = dataset_ops.DatasetV2.zip((dataset_3, dataset_4))
+    with self.assertRaises(ValueError):
+      py_builtins.all_(dataset_mixed)
+
+  def test_sorted(self):
+    self.assertListEqual(py_builtins.sorted_([2, 3, 1]), [1, 2, 3])
+    self.assertListEqual(
+        py_builtins.sorted_([2, 3, 1], key=lambda x: -x), [3, 2, 1])
+    self.assertListEqual(
+        py_builtins.sorted_([2, 3, 1], reverse=True), [3, 2, 1])
+    self.assertListEqual(
+        py_builtins.sorted_([2, 3, 1], key=lambda x: -x, reverse=True),
+        [1, 2, 3])
+    self.assertAllEqual(
+        py_builtins.sorted_([[4, 3], [2, 1]], key=lambda x: sum(x)),
+        [[2, 1], [4, 3]])
+
+  def test_sorted_tensor(self):
+    iterable_1 = constant_op.constant([2, 3, 1])
+    self.assertListEqual(
+        list(self.evaluate(py_builtins.sorted_(iterable_1))), [1, 2, 3])
+    self.assertListEqual(
+        list(self.evaluate(py_builtins.sorted_(iterable_1, key=lambda x: -x))),
+        [3, 2, 1])
+    self.assertListEqual(
+        list(self.evaluate(py_builtins.sorted_(iterable_1, reverse=True))),
+        [3, 2, 1])
+    self.assertListEqual(
+        list(
+            self.evaluate(
+                py_builtins.sorted_(iterable_1, key=lambda x: -x,
+                                    reverse=True))), [1, 2, 3])
+
+    iterable_2 = constant_op.constant([[4, 3], [2, 1]])
+    with self.assertRaises(ValueError):
+      py_builtins.sorted_(iterable_2)
+    with self.assertRaises(ValueError):
+      py_builtins.sorted_(iterable_2, key=lambda x: -x)
+    self.assertAllEqual(
+        list(
+            self.evaluate(
+                py_builtins.sorted_(
+                    iterable_2, key=lambda x: math_ops.reduce_sum(x)))),
+        [[2, 1], [4, 3]])
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/autograph/operators/slices.py

@@ -0,0 +1,142 @@
+# 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.
+# ==============================================================================
+"""Operators specific to slicing operations."""
+
+import collections
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor_util
+from tensorflow.python.ops import gen_array_ops
+from tensorflow.python.ops import gen_string_ops
+from tensorflow.python.ops import list_ops
+from tensorflow.python.ops import tensor_array_ops
+
+
+# TODO(mdan): Support extended slices.
+
+
+class GetItemOpts(collections.namedtuple('GetItemOpts', ('element_dtype',))):
+  pass
+
+
+def get_item(target, i, opts):
+  """The slice read operator (i.e. __getitem__).
+
+  Note: it is unspecified whether target will be mutated or not. In general,
+  if target is mutable (like Python lists), it will be mutated.
+
+  Args:
+    target: An entity that supports getitem semantics.
+    i: Index to read from.
+    opts: A GetItemOpts object.
+
+  Returns:
+    The read element.
+
+  Raises:
+    ValueError: if target is not of a supported type.
+  """
+  assert isinstance(opts, GetItemOpts)
+
+  if isinstance(target, tensor_array_ops.TensorArray):
+    return _tf_tensorarray_get_item(target, i)
+  elif tensor_util.is_tf_type(target):
+    if target.dtype == dtypes.variant:
+      return _tf_tensor_list_get_item(target, i, opts)
+    elif target.dtype == dtypes.string and target.shape.ndims == 0:
+      return _tf_tensor_string_get_item(target, i)
+    else:
+      return _tf_tensor_get_item(target, i)
+  else:
+    return _py_get_item(target, i)
+
+
+def _tf_tensorarray_get_item(target, i):
+  """Overload of get_item that stages a TensorArray read."""
+  return target.read(i)
+
+
+def _tf_tensor_list_get_item(target, i, opts):
+  """Overload of get_item that stages a Tensor list read."""
+  if opts.element_dtype is None:
+    raise ValueError('cannot retrieve from a list without knowing its '
+                     'element type; use set_element_type to annotate it')
+  x = list_ops.tensor_list_get_item(target, i, element_dtype=opts.element_dtype)
+  return x
+
+
+def _tf_tensor_get_item(target, i):
+  """Overload of get_item that stages a Tensor (not Tensor list) read."""
+  return target[i]
+
+
+def _tf_tensor_string_get_item(target, i):
+  """Overload of get_item that stages a Tensor string read."""
+  x = gen_string_ops.substr(target, i, 1)
+  return x
+
+
+def _py_get_item(target, i):
+  """Overload of get_item that executes a Python list modification."""
+  return target[i]
+
+
+def set_item(target, i, x):
+  """The slice write operator (i.e. __setitem__).
+
+  Note: it is unspecified whether target will be mutated or not. In general,
+  if target is mutable (like Python lists), it will be mutated.
+
+  Args:
+    target: An entity that supports setitem semantics.
+    i: Index to modify.
+    x: The new element value.
+
+  Returns:
+    Same as target, after the update was performed.
+
+  Raises:
+    ValueError: if target is not of a supported type.
+  """
+  if isinstance(target, tensor_array_ops.TensorArray):
+    return _tf_tensorarray_set_item(target, i, x)
+  elif tensor_util.is_tf_type(target):
+    if target.dtype == dtypes.variant:
+      return _tf_tensor_list_set_item(target, i, x)
+    else:
+      return _tf_tensor_set_item(target, i, x)
+  else:
+    return _py_set_item(target, i, x)
+
+
+def _tf_tensorarray_set_item(target, i, x):
+  """Overload of set_item that stages a TensorArray write."""
+  return target.write(i, x)
+
+
+def _tf_tensor_list_set_item(target, i, x):
+  """Overload of set_item that stages a Tensor list update."""
+  return list_ops.tensor_list_set_item(target, i, x)
+
+
+def _tf_tensor_set_item(target, i, x):
+  """Overload of set_item that stages a Tensor scatter update."""
+  return gen_array_ops.tensor_scatter_update(target, ((i,),), (x,))
+
+
+def _py_set_item(target, i, x):
+  """Overload of set_item that executes a Python list modification."""
+  target[i] = x
+  return target

tensorflow/python/autograph/operators/slices_test.py

@@ -0,0 +1,62 @@
+# 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 slices module."""
+
+from tensorflow.python.autograph.operators import slices
+from tensorflow.python.framework import constant_op
+from tensorflow.python.ops import list_ops
+from tensorflow.python.platform import test
+
+
+class SlicesTest(test.TestCase):
+
+  def test_set_item_tensor_list(self):
+    initial_list = constant_op.constant([[1, 2], [3, 4]])
+    elem_shape = constant_op.constant([2])
+    l = list_ops.tensor_list_from_tensor(initial_list, element_shape=elem_shape)
+    l = slices.set_item(l, 0, [5, 6])
+
+    with self.cached_session() as sess:
+      t = list_ops.tensor_list_stack(l, element_dtype=initial_list.dtype)
+      self.assertAllEqual(self.evaluate(t), [[5, 6], [3, 4]])
+
+  def test_get_item_tensor_list(self):
+    initial_list = constant_op.constant([[1, 2], [3, 4]])
+    elem_shape = constant_op.constant([2])
+    l = list_ops.tensor_list_from_tensor(initial_list, element_shape=elem_shape)
+    t = slices.get_item(
+        l, 1, slices.GetItemOpts(element_dtype=initial_list.dtype))
+
+    with self.cached_session() as sess:
+      self.assertAllEqual(self.evaluate(t), [3, 4])
+
+  def test_get_item_tensor_string(self):
+    initial_str = constant_op.constant('abcd')
+    t = slices.get_item(initial_str, 1,
+                        slices.GetItemOpts(element_dtype=initial_str.dtype))
+
+    with self.cached_session() as sess:
+      self.assertEqual(self.evaluate(t), b'b')
+
+    initial_list_str = constant_op.constant(['abcd', 'bcde'])
+    t = slices.get_item(initial_list_str, 1,
+                        slices.GetItemOpts(element_dtype=initial_str.dtype))
+
+    with self.cached_session() as sess:
+      self.assertEqual(self.evaluate(t), b'bcde')
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/autograph/operators/variables.py

@@ -0,0 +1,104 @@
+# Copyright 2019 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.
+# ==============================================================================
+"""Utilities used to capture Python idioms."""
+
+
+def ld(v):
+  """Load variable operator."""
+  if isinstance(v, Undefined):
+    return v.read()
+  return v
+
+
+def ldu(load_v, name):
+  """Load variable operator that returns Undefined when failing to evaluate.
+
+  Note: the name ("load or return undefined") is abbreviated to minimize
+  the amount of clutter in generated code.
+
+  This variant of `ld` is useful when loading symbols that may be undefined at
+  runtime, such as composite symbols, and whether they are defined or not cannot
+  be determined statically. For example `d['a']` is undefined when `d` is an
+  empty dict.
+
+  Args:
+    load_v: Lambda that executes the actual read.
+    name: Human-readable name of the symbol being read.
+  Returns:
+    Either the value of the symbol, or Undefined, if the symbol is not fully
+    defined.
+  """
+  try:
+    # TODO(mdan): Use locals()/globals() here.
+    return load_v()
+  except (KeyError, AttributeError, NameError):
+    return Undefined(name)
+
+
+class Undefined(object):
+  """Represents an undefined symbol in Python.
+
+  This is used to reify undefined symbols, which is required to use the
+  functional form of loops.
+  Example:
+
+    while n > 0:
+      n = n - 1
+      s = n
+    return s  # Runtime error if n == 0
+
+  This is valid Python code and will not result in an error as long as n
+  is positive. The use of this class is to stay as close to Python semantics
+  as possible for staged code of this nature.
+
+  Converted version of the above showing the possible usage of this class:
+
+    s = Undefined('s')
+    init_state = (s,)
+    s = while_loop(cond, body, init_state)
+    return s  # s is an instance of Undefined if the loop never runs
+
+  Attributes:
+    symbol_name: Text, identifier for the undefined symbol
+  """
+
+  __slots__ = ('symbol_name',)
+
+  def __init__(self, symbol_name):
+    self.symbol_name = symbol_name
+
+  def read(self):
+    raise UnboundLocalError("'{}' is used before assignment".format(
+        self.symbol_name))
+
+  def __repr__(self):
+    return self.symbol_name
+
+  def __getattribute__(self, name):
+    try:
+      # If it's an existing attribute, return it.
+      return object.__getattribute__(self, name)
+    except AttributeError:
+      # Otherwise return Undefined.
+      return self
+
+  def __getitem__(self, i):
+    return self
+
+
+# TODO(mdan): Refactor as a RetVal object, aggregating the value and do_return.
+class UndefinedReturnValue(object):
+  """Represents a return value that is undefined."""
+  pass

tensorflow/python/autograph/operators/variables_test.py

@@ -0,0 +1,51 @@
+# Copyright 2019 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 python_lang_utils module."""
+
+from tensorflow.python.autograph.operators import variables
+from tensorflow.python.platform import test
+
+
+class SpecialValuesTest(test.TestCase):
+
+  def test_undefined(self):
+    undefined_symbol = variables.Undefined('name')
+    undefined_symbol2 = variables.Undefined('name')
+
+    self.assertEqual(undefined_symbol.symbol_name, 'name')
+    self.assertEqual(undefined_symbol2.symbol_name, 'name')
+    self.assertNotEqual(undefined_symbol, undefined_symbol2)
+
+  def test_undefined_operations(self):
+    undefined_symbol = variables.Undefined('name')
+
+    self.assertIsInstance(undefined_symbol.foo, variables.Undefined)
+    self.assertIsInstance(undefined_symbol[0], variables.Undefined)
+    self.assertNotIsInstance(undefined_symbol.__class__, variables.Undefined)
+
+  def test_read(self):
+    self.assertEqual(variables.ld(1), 1)
+    o = object()
+    self.assertEqual(variables.ld(o), o)
+
+    self.assertIsNone(variables.ld(None))
+
+  def test_read_undefined(self):
+    with self.assertRaisesRegex(UnboundLocalError, 'used before assignment'):
+      variables.ld(variables.Undefined('a'))
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/autograph/pyct/BUILD

@@ -0,0 +1,416 @@
+load("@rules_shell//shell:sh_test.bzl", "sh_test")
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:tensorflow.bzl", "py_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "templates",
+    srcs = ["templates.py"],
+    strict_deps = True,
+    visibility = ["//visibility:public"],
+    deps = [
+        ":anno",
+        ":ast_util",
+        ":parser",
+        ":qual_names",
+        "@pypi//gast",
+    ],
+)
+
+py_library(
+    name = "transpiler",
+    srcs = ["transpiler.py"],
+    strict_deps = True,
+    visibility = ["//visibility:public"],
+    deps = [
+        ":cache",
+        ":inspect_utils",
+        ":loader",
+        ":naming",
+        ":origin_info",
+        ":parser",
+        ":templates",
+        ":transformer",
+        "//tensorflow/python/autograph/utils:ag_logging",
+        "@pypi//gast",
+    ],
+)
+
+py_library(
+    name = "ast_util",
+    srcs = ["ast_util.py"],
+    strict_deps = True,
+    visibility = ["//visibility:public"],
+    deps = [
+        ":anno",
+        ":parser",
+        ":qual_names",
+        "@pypi//gast",
+    ],
+)
+
+py_library(
+    name = "loader",
+    srcs = ["loader.py"],
+    strict_deps = True,
+    visibility = ["//visibility:public"],
+    deps = [
+        ":origin_info",
+        ":parser",
+    ],
+)
+
+py_library(
+    name = "gast_util",
+    srcs = ["gast_util.py"],
+    strict_deps = True,
+    visibility = ["//visibility:public"],
+    deps = ["@pypi//gast"],
+)
+
+py_library(
+    name = "__init__",
+    srcs = ["__init__.py"],
+    strict_deps = True,
+    visibility = ["//visibility:public"],
+)
+
+py_library(
+    name = "parser",
+    srcs = ["parser.py"],
+    strict_deps = True,
+    visibility = ["//visibility:public"],
+    deps = [
+        ":errors",
+        ":inspect_utils",
+        "//tensorflow/python/util:tf_inspect",
+        "@pypi//gast",
+    ],
+)
+
+py_library(
+    name = "naming",
+    srcs = ["naming.py"],
+    strict_deps = True,
+    visibility = ["//visibility:public"],
+    deps = [":qual_names"],
+)
+
+py_library(
+    name = "inspect_utils",
+    srcs = ["inspect_utils.py"],
+    strict_deps = True,
+    visibility = ["//visibility:public"],
+    deps = ["//tensorflow/python/util:tf_inspect"],
+)
+
+py_library(
+    name = "origin_info",
+    srcs = ["origin_info.py"],
+    strict_deps = True,
+    visibility = ["//visibility:public"],
+    deps = [
+        ":anno",
+        ":ast_util",
+        ":parser",
+        ":pretty_printer",
+        "//tensorflow/python/util:tf_inspect",
+        "@pypi//gast",
+    ],
+)
+
+py_library(
+    name = "anno",
+    srcs = ["anno.py"],
+    strict_deps = True,
+    visibility = ["//visibility:public"],
+    deps = ["@pypi//gast"],
+)
+
+py_library(
+    name = "errors",
+    srcs = ["errors.py"],
+    strict_deps = True,
+    visibility = ["//visibility:public"],
+)
+
+py_library(
+    name = "transformer",
+    srcs = ["transformer.py"],
+    strict_deps = True,
+    visibility = ["//visibility:public"],
+    deps = [
+        ":anno",
+        ":parser",
+        ":pretty_printer",
+        ":templates",
+        "@pypi//gast",
+    ],
+)
+
+py_library(
+    name = "qual_names",
+    srcs = ["qual_names.py"],
+    strict_deps = True,
+    visibility = ["//visibility:public"],
+    deps = [
+        ":anno",
+        ":parser",
+        "@pypi//gast",
+    ],
+)
+
+py_library(
+    name = "cfg",
+    srcs = ["cfg.py"],
+    strict_deps = True,
+    visibility = ["//visibility:public"],
+    deps = [
+        ":anno",
+        "@pypi//gast",
+    ],
+)
+
+py_library(
+    name = "error_utils",
+    srcs = ["error_utils.py"],
+    strict_deps = True,
+    visibility = ["//visibility:public"],
+    deps = [
+        ":origin_info",
+        "//tensorflow/python/util:traceback_utils",
+    ],
+)
+
+py_library(
+    name = "cache",
+    srcs = ["cache.py"],
+    strict_deps = True,
+    visibility = ["//visibility:public"],
+)
+
+py_library(
+    name = "pretty_printer",
+    srcs = ["pretty_printer.py"],
+    strict_deps = True,
+    visibility = ["//visibility:public"],
+    deps = [
+        "@pypi//gast",
+        "@pypi//termcolor",
+    ],
+)
+
+py_test(
+    name = "anno_test",
+    srcs = ["anno_test.py"],
+    strict_deps = True,
+    deps = [
+        ":anno",
+        #internal proto upb dep
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_test(
+    name = "ast_util_test",
+    srcs = ["ast_util_test.py"],
+    strict_deps = True,
+    deps = [
+        ":anno",
+        ":ast_util",
+        ":loader",
+        ":parser",
+        ":pretty_printer",
+        ":qual_names",
+        "@pypi//gast",
+        #internal proto upb dep
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_test(
+    name = "cache_test",
+    srcs = ["cache_test.py"],
+    strict_deps = True,
+    deps = [
+        ":cache",
+        #internal proto upb dep
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_test(
+    name = "cfg_test",
+    srcs = ["cfg_test.py"],
+    strict_deps = True,
+    deps = [
+        ":cfg",
+        ":parser",
+        "@pypi//gast",
+        #internal proto upb dep
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_test(
+    name = "loader_test",
+    srcs = ["loader_test.py"],
+    strict_deps = True,
+    deps = [
+        ":ast_util",
+        ":loader",
+        ":parser",
+        ":pretty_printer",
+        "@pypi//gast",
+        #internal proto upb dep
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/util:tf_inspect",
+    ],
+)
+
+py_test(
+    name = "error_utils_test",
+    srcs = ["error_utils_test.py"],
+    strict_deps = True,
+    deps = [
+        ":error_utils",
+        ":origin_info",
+        #internal proto upb dep
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_test(
+    name = "inspect_utils_test",
+    srcs = ["inspect_utils_test.py"],
+    strict_deps = True,
+    deps = [
+        ":inspect_utils",
+        #internal proto upb dep
+        "//tensorflow/python/autograph/pyct/testing:basic_definitions",
+        "//tensorflow/python/autograph/pyct/testing:decorators",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/lib:__init__",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+sh_test(
+    name = "inspect_utils_test_par",
+    srcs = ["inspect_utils_test.sh"],
+    tags = ["no_oss"],
+)
+
+py_test(
+    name = "naming_test",
+    srcs = ["naming_test.py"],
+    strict_deps = True,
+    deps = [
+        ":naming",
+        #internal proto upb dep
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_test(
+    name = "origin_info_test",
+    srcs = ["origin_info_test.py"],
+    strict_deps = True,
+    deps = [
+        ":anno",
+        ":inspect_utils",
+        ":origin_info",
+        ":parser",
+        #internal proto upb dep
+        "//tensorflow/python/autograph/pyct/testing:basic_definitions",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/util:tf_inspect",
+    ],
+)
+
+py_test(
+    name = "parser_test",
+    srcs = ["parser_test.py"],
+    strict_deps = True,
+    deps = [
+        ":ast_util",
+        ":errors",
+        ":parser",
+        ":pretty_printer",
+        "@pypi//gast",
+        #internal proto upb dep
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_test(
+    name = "pretty_printer_test",
+    srcs = ["pretty_printer_test.py"],
+    strict_deps = True,
+    deps = [
+        ":pretty_printer",
+        #internal proto upb dep
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_test(
+    name = "qual_names_test",
+    srcs = ["qual_names_test.py"],
+    strict_deps = True,
+    deps = [
+        ":anno",
+        ":parser",
+        ":qual_names",
+        #internal proto upb dep
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_test(
+    name = "templates_test",
+    srcs = ["templates_test.py"],
+    strict_deps = True,
+    deps = [
+        ":loader",
+        ":parser",
+        ":qual_names",
+        ":templates",
+        "@absl_py//absl/testing:parameterized",
+        "@pypi//gast",
+        #internal proto upb dep
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_test(
+    name = "transformer_test",
+    srcs = ["transformer_test.py"],
+    strict_deps = True,
+    deps = [
+        ":anno",
+        ":origin_info",
+        ":parser",
+        ":transformer",
+        "@pypi//gast",
+        #internal proto upb dep
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)
+
+py_test(
+    name = "transpiler_test",
+    srcs = ["transpiler_test.py"],
+    strict_deps = True,
+    deps = [
+        ":transformer",
+        ":transpiler",
+        "@pypi//gast",
+        #internal proto upb dep
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)

tensorflow/python/autograph/pyct/__init__.py

@@ -0,0 +1,16 @@
+# 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.
+# ==============================================================================
+"""Python source code transformation library."""
+

tensorflow/python/autograph/pyct/anno.py

@@ -0,0 +1,174 @@
+# 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.
+# ==============================================================================
+"""AST node annotation support.
+
+Adapted from Tangent.
+"""
+
+import enum
+
+# pylint:disable=g-bad-import-order
+
+import gast
+# pylint:enable=g-bad-import-order
+
+
+# TODO(mdan): Shorten the names.
+# These names are heavily used, and anno.blaa
+# TODO(mdan): Replace the attr-dict mechanism with a more typed solution.
+
+
+class NoValue(enum.Enum):
+  """Base class for different types of AST annotations."""
+
+  def of(self, node, default=None):
+    return getanno(node, self, default=default)
+
+  def add_to(self, node, value):
+    setanno(node, self, value)
+
+  def exists(self, node):
+    return hasanno(node, self)
+
+  def __repr__(self):
+    return str(self.name)
+
+
+class Basic(NoValue):
+  """Container for basic annotation keys.
+
+  The enum values are used strictly for documentation purposes.
+  """
+
+  QN = 'Qualified name, as it appeared in the code. See qual_names.py.'
+  SKIP_PROCESSING = (
+      'This node should be preserved as is and not processed any further.')
+  INDENT_BLOCK_REMAINDER = (
+      'When a node is annotated with this, the remainder of the block should'
+      ' be indented below it. The annotation contains a tuple'
+      ' (new_body, name_map), where `new_body` is the new indented block and'
+      ' `name_map` allows renaming symbols.')
+  ORIGIN = ('Information about the source code that converted code originated'
+            ' from. See origin_information.py.')
+  DIRECTIVES = ('User directives associated with a statement or a variable.'
+                ' Typically, they affect the immediately-enclosing statement.')
+
+  EXTRA_LOOP_TEST = (
+      'A special annotation containing additional test code to be executed in'
+      ' for loops.')
+
+
+class Static(NoValue):
+  """Container for static analysis annotation keys.
+
+  The enum values are used strictly for documentation purposes.
+  """
+
+  # Symbols
+  # These flags are boolean.
+  IS_PARAM = 'Symbol is a parameter to the function being analyzed.'
+
+  # Scopes
+  # Scopes are represented by objects of type activity.Scope.
+  SCOPE = 'The scope for the annotated node. See activity.py.'
+  # TODO(mdan): Drop these in favor of accessing the child's SCOPE.
+  ARGS_SCOPE = 'The scope for the argument list of a function call.'
+  COND_SCOPE = 'The scope for the test node of a conditional statement.'
+  BODY_SCOPE = (
+      'The scope for the main body of a statement (True branch for if '
+      'statements, main body for loops).')
+  ORELSE_SCOPE = (
+      'The scope for the orelse body of a statement (False branch for if '
+      'statements, orelse body for loops).')
+
+  # Static analysis annotations.
+  DEFINITIONS = (
+      'Reaching definition information. See reaching_definitions.py.')
+  ORIG_DEFINITIONS = (
+      'The value of DEFINITIONS that applied to the original code before any'
+      ' conversion.')
+  DEFINED_FNS_IN = (
+      'Local function definitions that may exist when exiting the node. See'
+      ' reaching_fndefs.py')
+  DEFINED_VARS_IN = (
+      'Symbols defined when entering the node. See reaching_definitions.py.')
+  LIVE_VARS_OUT = ('Symbols live when exiting the node. See liveness.py.')
+  LIVE_VARS_IN = ('Symbols live when entering the node. See liveness.py.')
+  TYPES = 'Static type information. See type_inference.py.'
+  CLOSURE_TYPES = 'Types of closure symbols at each detected call site.'
+  VALUE = 'Static value information. See type_inference.py.'
+
+
+FAIL = object()
+
+
+def keys(node, field_name='___pyct_anno'):
+  if not hasattr(node, field_name):
+    return frozenset()
+  return frozenset(getattr(node, field_name).keys())
+
+
+def getanno(node, key, default=FAIL, field_name='___pyct_anno'):
+  if (default is FAIL or (hasattr(node, field_name) and
+                          (key in getattr(node, field_name)))):
+    return getattr(node, field_name)[key]
+  return default
+
+
+def hasanno(node, key, field_name='___pyct_anno'):
+  return hasattr(node, field_name) and key in getattr(node, field_name)
+
+
+def setanno(node, key, value, field_name='___pyct_anno'):
+  annotations = getattr(node, field_name, {})
+  setattr(node, field_name, annotations)
+  annotations[key] = value
+
+  # So that the annotations survive gast_to_ast() and ast_to_gast()
+  if field_name not in node._fields:
+    node._fields += (field_name,)
+
+
+def delanno(node, key, field_name='___pyct_anno'):
+  annotations = getattr(node, field_name)
+  del annotations[key]
+  if not annotations:
+    delattr(node, field_name)
+    node._fields = tuple(f for f in node._fields if f != field_name)
+
+
+def copyanno(from_node, to_node, key, field_name='___pyct_anno'):
+  if hasanno(from_node, key, field_name=field_name):
+    setanno(
+        to_node,
+        key,
+        getanno(from_node, key, field_name=field_name),
+        field_name=field_name)
+
+
+def dup(node, copy_map, field_name='___pyct_anno'):
+  """Recursively copies annotations in an AST tree.
+
+  Args:
+    node: ast.AST
+    copy_map: Dict[Hashable, Hashable], maps a source anno key to a destination
+        key. All annotations with the source key will be copied to identical
+        annotations with the destination key.
+    field_name: str
+  """
+  for n in gast.walk(node):
+    for k in copy_map:
+      if hasanno(n, k, field_name):
+        setanno(n, copy_map[k], getanno(n, k, field_name), field_name)