xref: /aosp_15_r20/external/federated-compute/fcp/tensorflow/tensor_name.py (revision 14675a029014e728ec732f129a32e299b2da0601)
1# Copyright 2022 Google LLC
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"""Provides the `tensor_name` operation.
15
16This wraps the generated op and ensures that necessary shared libraries
17are loaded.
18"""
19
20import tensorflow as tf
21
22from fcp.tensorflow import gen_tensor_name_py
23
24_tensor_name_so = tf.load_op_library(
25    tf.compat.v1.resource_loader.get_path_to_datafile('./_tensor_name_op.so'))
26
27
28def tensor_name(tensor):
29  """Returns the final graph name of a tensor as a string tensor."""
30  if not tf.is_tensor(tensor):
31    raise TypeError('`tensor_name` expected a tensor, found object of type '
32                    f'{type(tensor)}.')
33  return gen_tensor_name_py.tensor_name(input_tensor=tensor)
34