181 lines
6.9 KiB
Python
181 lines
6.9 KiB
Python
# Copyright 2018 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.
|
|
# ==============================================================================
|
|
"""Ops to convert between RaggedTensors and other tensor types."""
|
|
|
|
from tensorflow.python.framework import dtypes
|
|
from tensorflow.python.framework import indexed_slices
|
|
from tensorflow.python.framework import ops
|
|
from tensorflow.python.ops import array_ops
|
|
from tensorflow.python.ops import gen_ragged_conversion_ops
|
|
from tensorflow.python.ops import math_ops
|
|
from tensorflow.python.ops.ragged import ragged_tensor
|
|
|
|
|
|
def from_tensor(tensor,
|
|
lengths=None,
|
|
padding=None,
|
|
ragged_rank=1,
|
|
row_splits_dtype=dtypes.int64,
|
|
name=None):
|
|
if ragged_tensor.is_ragged(tensor):
|
|
return tensor
|
|
else:
|
|
return ragged_tensor.RaggedTensor.from_tensor(
|
|
tensor,
|
|
lengths=lengths,
|
|
padding=padding,
|
|
ragged_rank=ragged_rank,
|
|
row_splits_dtype=row_splits_dtype,
|
|
name=name)
|
|
|
|
|
|
def to_tensor(rt_input, default_value=None, name=None):
|
|
if ragged_tensor.is_ragged(rt_input):
|
|
return rt_input.to_tensor(default_value, name)
|
|
else:
|
|
return rt_input
|
|
|
|
|
|
def ragged_to_dense(rt_input, default_value=None, shape=None):
|
|
"""Create a dense tensor from a ragged tensor."""
|
|
return rt_input.to_tensor(default_value=default_value, shape=shape)
|
|
|
|
|
|
@ops.RegisterGradient("RaggedTensorToTensor")
|
|
def _ragged_tensor_to_tensor_grad(op, grad):
|
|
"""Gradient for RaggedToTensor op."""
|
|
# Extract inputs from the op.
|
|
flat_values = op.inputs[1]
|
|
default_value = op.inputs[2]
|
|
row_partition_tensors = op.inputs[3:]
|
|
row_partition_types = op.get_attr("row_partition_types")
|
|
flat_value_shape = array_ops.shape(flat_values)
|
|
ragged_rank = sum(
|
|
1 for typ in row_partition_types if typ != b"FIRST_DIM_SIZE")
|
|
|
|
# Create two tensors that correspond 1:1 with grad (and op.output):
|
|
# * indices[i1...iN] is the index in `flat_values` of the value used to
|
|
# populate output[i1...iN] (if the value came from `flat_values`) or
|
|
# -1 (if the value came from `default_value`).
|
|
# * mask[i1...iN] is true if output[i1...iN] came from `flat_values`, or
|
|
# false if it came from `default_value`.
|
|
indices = gen_ragged_conversion_ops.ragged_tensor_to_tensor(
|
|
shape=array_ops.shape(grad)[:1 + ragged_rank],
|
|
values=math_ops.range(flat_value_shape[0]),
|
|
default_value=-1,
|
|
row_partition_types=row_partition_types,
|
|
row_partition_tensors=row_partition_tensors)
|
|
mask = math_ops.not_equal(indices, -1)
|
|
|
|
# Select out the gradients & indices that came from `flat_values`, and use
|
|
# those to construct the gradient for `flat_values` (as an IndexedSlices).
|
|
values_grad = indexed_slices.IndexedSlices(
|
|
values=array_ops.boolean_mask(grad, mask),
|
|
indices=array_ops.boolean_mask(indices, mask),
|
|
dense_shape=flat_value_shape)
|
|
|
|
# Select out the gradients that came from `default_value`, and sum them to
|
|
# get the gradient for the default. Note that the default_value may have
|
|
# been broadcast as part of the RaggedTensorToTensor operation, so we also
|
|
# need to reduce any dimensions that might have been broadcast.
|
|
default_grads = array_ops.boolean_mask(grad, ~mask)
|
|
dims_to_reduce = math_ops.range(
|
|
array_ops.rank(default_grads) -
|
|
_rank_ignoring_leading_dims_with_size_1(default_value))
|
|
default_grad = math_ops.reduce_sum(default_grads, axis=dims_to_reduce)
|
|
|
|
# Restore any leading dims with size one.
|
|
default_grad = array_ops.reshape(default_grad, array_ops.shape(default_value))
|
|
|
|
return ([None, values_grad, default_grad] +
|
|
[None for _ in row_partition_tensors])
|
|
|
|
|
|
def _rank_ignoring_leading_dims_with_size_1(value):
|
|
"""Returns `rank(value)`, ignoring any leading dimensions with size 1."""
|
|
# Compute the result using static shape, if possible.
|
|
if value.shape.rank is not None:
|
|
ndims = value.shape.rank
|
|
for dim in value.shape.dims:
|
|
if dim.value == 1:
|
|
ndims -= 1
|
|
elif dim.value is None:
|
|
ndims = None # Can't compute the result using static shape.
|
|
break
|
|
else:
|
|
break
|
|
if ndims is not None:
|
|
return ndims
|
|
|
|
# Otherwise, we need to compute the result dynamically. The math we use to
|
|
# do this is a bit round-about, so here's an example to illustrate:
|
|
# shape = [1, 1, 3, 5, 1, 4] # shape(value)
|
|
# dim_is_one = [1, 1, 0, 0, 1, 0] # equal(shape, 1)
|
|
# leading_ones = [1, 1, 0, 0, 0, 0] # cumprod(dim_is_one)
|
|
# num_leading_ones = 2 # reduce_sum(leading_ones)
|
|
# result = 4 # rank(value) - num_leading_ones
|
|
shape = array_ops.shape(value)
|
|
dim_is_one = math_ops.cast(math_ops.equal(shape, 1), dtypes.int32)
|
|
leading_ones = math_ops.cumprod(dim_is_one)
|
|
num_leading_ones = math_ops.reduce_sum(leading_ones)
|
|
return array_ops.rank(value) - num_leading_ones
|
|
|
|
|
|
def to_sparse(rt_input, name=None):
|
|
return rt_input.to_sparse(name)
|
|
|
|
|
|
def from_sparse(st_input, name=None):
|
|
return ragged_tensor.RaggedTensor.from_sparse(st_input, name)
|
|
|
|
|
|
@ops.RegisterGradient("RaggedTensorFromVariant")
|
|
def _ragged_tensor_from_variant_grad(op, *grads):
|
|
"""Gradient for RaggedTensorFromVariant op."""
|
|
|
|
variant_rank = op.inputs[0].shape.rank
|
|
if variant_rank == 0:
|
|
batched_input = False
|
|
elif variant_rank == 1:
|
|
batched_input = True
|
|
elif variant_rank is None:
|
|
batched_input = (op.get_attr("output_ragged_rank") > 0)
|
|
else:
|
|
# TODO(edloper): Add a batch_dims argument to RaggedTensorToVariant, so
|
|
# we can support this.
|
|
raise ValueError("Unable to compute gradient: RaggedTensorToVariant "
|
|
"can currently only generate 0D or 1D output.")
|
|
return [
|
|
gen_ragged_conversion_ops.ragged_tensor_to_variant(
|
|
rt_nested_splits=op.outputs[:-1],
|
|
rt_dense_values=grads[-1],
|
|
batched_input=batched_input)
|
|
]
|
|
|
|
|
|
@ops.RegisterGradient("RaggedTensorToVariant")
|
|
def _ragged_tensor_to_variant_grad(op, encoded_ragged_grad):
|
|
"""Gradient for RaggedTensorToVariant op."""
|
|
dense_values = op.inputs[-1]
|
|
ragged_rank = len(op.inputs) - 1
|
|
row_splits = 0 if ragged_rank == 0 else op.inputs[0]
|
|
values_grad = gen_ragged_conversion_ops.ragged_tensor_to_variant_gradient(
|
|
encoded_ragged_grad=encoded_ragged_grad,
|
|
row_splits=row_splits,
|
|
dense_values_shape=array_ops.shape(dense_values),
|
|
Tvalues=op.inputs[-1].dtype)
|
|
result = [None] * ragged_rank + [values_grad]
|
|
return result
|