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"""Dynamic shape for structured Tensors.""" 16 17from tensorflow.python.framework import ops 18from tensorflow.python.framework import tensor_shape 19from tensorflow.python.ops import array_ops 20from tensorflow.python.ops.ragged import dynamic_ragged_shape 21from tensorflow.python.ops.structured.structured_tensor import _find_shape_dtype 22 23 24# pylint:disable=protected-access 25def _dynamic_ragged_shape_init(fields, shape, nrows, row_partitions): 26 """Produce a DynamicRaggedShape for StructuredTensor.""" 27 assert isinstance(fields, dict), fields 28 assert isinstance(shape, tensor_shape.TensorShape), shape 29 assert nrows is None or isinstance(nrows, ops.Tensor), nrows 30 assert isinstance(row_partitions, tuple), row_partitions 31 32 rank = shape.rank 33 if rank is None: 34 raise TypeError("StructuredTensor's shape must have known rank.") 35 36 # TODO(martinz): figure out whether to validate. 37 dtype = _find_shape_dtype(fields, nrows, row_partitions) 38 if rank == 0: 39 return dynamic_ragged_shape.DynamicRaggedShape._from_inner_shape( 40 array_ops.zeros((0,), dtype=dtype)) 41 42 if rank == 1: 43 alt_value = shape[0] 44 if isinstance(alt_value, tensor_shape.Dimension): 45 alt_value = alt_value.value 46 if alt_value is not None: 47 nrows = alt_value 48 return dynamic_ragged_shape.DynamicRaggedShape._from_inner_shape( 49 [nrows], dtype=dtype) 50 51 return dynamic_ragged_shape.DynamicRaggedShape.from_row_partitions( 52 row_partitions, dtype=dtype) 53