xref: /aosp_15_r20/external/tensorflow/tensorflow/python/trackable/converter.py (revision b6fb3261f9314811a0f4371741dbb8839866f948)
1# Copyright 2022 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"""Util for converting a Python object to a Trackable."""
16
17
18from tensorflow.python.eager import function_saved_model_utils
19from tensorflow.python.framework import dtypes
20from tensorflow.python.framework import tensor_util
21from tensorflow.python.ops import resource_variable_ops
22from tensorflow.python.trackable import base
23from tensorflow.python.trackable import data_structures
24
25
26def convert_to_trackable(obj, parent=None):
27  """Converts `obj` to `Trackable`."""
28  if isinstance(obj, base.Trackable):
29    return obj
30  obj = data_structures.wrap_or_unwrap(obj)
31  if (tensor_util.is_tf_type(obj) and
32      obj.dtype not in (dtypes.variant, dtypes.resource) and
33      not resource_variable_ops.is_resource_variable(obj)):
34    return function_saved_model_utils.TrackableConstant(obj, parent)
35  if not isinstance(obj, base.Trackable):
36    raise ValueError(f"Cannot convert {obj} to Trackable.")
37  return obj
38