54 lines
2.4 KiB
Python
54 lines
2.4 KiB
Python
|
# Copyright 2022 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.
|
||
|
|
# ==============================================================================
|
||
|
|
"""The implementation of `tf.data.Dataset.from_sparse_tensor_slices`."""
|
||
|
|
|
||
|
|
from tensorflow.python.data.ops import dataset_ops
|
||
|
|
from tensorflow.python.framework import dtypes
|
||
|
|
from tensorflow.python.framework import sparse_tensor as sparse_tensor_lib
|
||
|
|
from tensorflow.python.framework import tensor_spec
|
||
|
|
from tensorflow.python.ops import gen_dataset_ops
|
||
|
|
|
||
|
|
|
||
|
|
def _from_sparse_tensor_slices(sparse_tensor): # pylint: disable=unused-private-name
|
||
|
|
return dataset_ops.DatasetV1Adapter(_SparseTensorSliceDataset(sparse_tensor))
|
||
|
|
|
||
|
|
|
||
|
|
class _SparseTensorSliceDataset(dataset_ops.DatasetSource):
|
||
|
|
"""A `Dataset` that splits a rank-N `tf.sparse.SparseTensor` into its rows."""
|
||
|
|
|
||
|
|
def __init__(self, sparse_tensor):
|
||
|
|
"""See `Dataset.from_sparse_tensor_slices()` for details."""
|
||
|
|
if not isinstance(sparse_tensor, sparse_tensor_lib.SparseTensor):
|
||
|
|
raise TypeError(f"Invalid `sparse_tensor`. `sparse_tensor` must be a "
|
||
|
|
f"`tf.sparse.SparseTensor`. Got {type(sparse_tensor)}.")
|
||
|
|
self._sparse_tensor = sparse_tensor
|
||
|
|
|
||
|
|
indices_shape = self._sparse_tensor.indices.get_shape()
|
||
|
|
shape_shape = self._sparse_tensor.dense_shape.get_shape()
|
||
|
|
rank = (indices_shape.dims[1] - 1).merge_with(shape_shape.dims[0] - 1)
|
||
|
|
self._structure = (tensor_spec.TensorSpec([None, rank], dtypes.int64),
|
||
|
|
tensor_spec.TensorSpec([None],
|
||
|
|
self._sparse_tensor.dtype),
|
||
|
|
tensor_spec.TensorSpec([rank], dtypes.int64))
|
||
|
|
|
||
|
|
variant_tensor = gen_dataset_ops.sparse_tensor_slice_dataset(
|
||
|
|
self._sparse_tensor.indices, self._sparse_tensor.values,
|
||
|
|
self._sparse_tensor.dense_shape)
|
||
|
|
super().__init__(variant_tensor)
|
||
|
|
|
||
|
|
@property
|
||
|
|
def element_spec(self):
|
||
|
|
return self._structure
|