wehub-resource-sync
75 lines
2.9 KiB
Python
75 lines
2.9 KiB
Python
# 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.
|
|
# ==============================================================================
|
|
"""Tensor shape utilities."""
|
|
from tensorflow.python.eager import context
|
|
from tensorflow.python.framework import dtypes
|
|
from tensorflow.python.framework import ops
|
|
from tensorflow.python.framework import tensor_shape
|
|
from tensorflow.python.framework import tensor_util
|
|
|
|
|
|
def shape_tensor(shape): # pylint: disable=invalid-name
|
|
"""Convert to an int32 or int64 tensor, defaulting to int32 if empty."""
|
|
dtype = None
|
|
if isinstance(shape, (tuple, list)):
|
|
if not shape:
|
|
dtype = dtypes.int32
|
|
else:
|
|
# If there are Dimension objects in the shape, unwrap them. This can be a
|
|
# problem if v1 and v2 TensorShape objects get mixed up in partial
|
|
# conversions, leading to shapes such as (1, 2, Dimension(5)), which are
|
|
# not convertible to Tensors because of mixed content.
|
|
shape = tuple(map(tensor_shape.dimension_value, shape))
|
|
return ops.convert_to_tensor(shape, dtype=dtype, name="shape")
|
|
|
|
|
|
# DO NOT USE: For testing only.
|
|
_ENABLE_MAYBE_SET_STATIC_SHAPE = True
|
|
|
|
|
|
def maybe_set_static_shape(tensor, shape): # pylint: disable=invalid-name
|
|
"""Sets the shape of `tensor` to the `shape`'s constant value, if inferrable.
|
|
|
|
This is a temporary workaround to fix shape inference across functional op
|
|
boundaries. E.g.
|
|
|
|
```python
|
|
shape = tf.constant([3])
|
|
@tf.function
|
|
def f():
|
|
u = tf.random_uniform(shape)
|
|
return u
|
|
```
|
|
|
|
If we were to rely solely on C++ shape inference, the shape of `u` inside
|
|
`f` would be unknown because C++ shape inference is not aware of the outer
|
|
graph and all it sees is a Placeholder node when backtracing the captured
|
|
tensor for `shape`. `maybe_set_static_shape` computes the static shape value
|
|
of `shape` by traversing the `FuncGraph` boundaries and sets the correct
|
|
shape.
|
|
|
|
A longer term solution would be to fix C++ shape inference.
|
|
|
|
Args:
|
|
tensor: A tensor.
|
|
shape: A shape tensor.
|
|
"""
|
|
if (_ENABLE_MAYBE_SET_STATIC_SHAPE and not context.executing_eagerly() and
|
|
ops.get_default_graph().building_function and
|
|
not tensor.shape.is_fully_defined() and tensor_util.is_tensor(shape)):
|
|
shape = shape_tensor(shape)
|
|
const_shape = tensor_util.constant_value_as_shape(shape)
|
|
tensor.set_shape(const_shape)
|