xref: /aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/tensorflow/tests/tf_saved_model/debug_info.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
16# RUN: %p/debug_info | FileCheck %s
17
18# pylint: disable=missing-docstring,line-too-long
19import tensorflow.compat.v2 as tf
20from tensorflow.compiler.mlir.tensorflow.tests.tf_saved_model import common
21
22
23class TestModule(tf.Module):
24
25  @tf.function(input_signature=[
26      tf.TensorSpec([], tf.float32),
27      tf.TensorSpec([], tf.float32)
28  ])
29  def some_function(self, x, y):
30    return x + y
31    # Basic check that the debug info file is being correctly saved and loaded.
32    #
33    # CHECK: "tf.AddV2"{{.*}}loc(#[[LOC:.*]])
34    # CHECK: #[[LOC]] = loc({{.*}}callsite("{{[^"]*}}/debug_info.py{{.*}}":{{[0-9]+}}:{{[0-9]+}}
35
36
37if __name__ == '__main__':
38  common.do_test(TestModule, show_debug_info=True)
39