xref: /aosp_15_r20/external/tensorflow/tensorflow/python/debug/lib/op_callbacks_common.py (revision b6fb3261f9314811a0f4371741dbb8839866f948)
1# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ==============================================================================
15"""Common utilities and settings used by tfdbg v2's op callbacks."""
16
17# The ops that are skipped by tfdbg v2's op callbacks.
18# They belong to TensorFlow's control flow ops (e.g., "Enter", "StatelessIf")
19# and ops that wrap nested tf.function calls.
20OP_CALLBACK_SKIP_OPS = (
21    # TODO(b/139668453): The following skipped ops are related to a limitation
22    # in the op callback.
23    b"Enter",
24    b"Exit",
25    b"Identity",
26    b"If",
27    b"LoopCond",
28    b"Merge",
29    b"NextIteration",
30    b"StatelessIf",
31    b"StatefulPartitionedCall",
32    b"Switch",
33    b"While",
34    # NOTE(b/154097452): On TPUs, debugger ops are colocated with RemoteCall
35    # ops. This exclusion prevents an error due to no OpKernel for those
36    # debugger ops.
37    b"RemoteCall",
38    # TPU-specific ops begin.
39    b"TPUReplicatedInput",
40    b"TPUReplicateMetadata",
41    b"TPUCompilationResult",
42    b"TPUReplicatedOutput",
43    b"ConfigureDistributedTPU",
44    # Other special ops used by TensorFlow internally.
45    b"DestroyResourceOp",
46)
47