1010 lines
42 KiB
Python
1010 lines
42 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.
|
|
# ==============================================================================
|
|
"""Tests for utilities working with arbitrarily nested structures."""
|
|
|
|
import collections
|
|
import dataclasses
|
|
import functools
|
|
|
|
from absl.testing import parameterized
|
|
import numpy as np
|
|
import wrapt
|
|
|
|
from tensorflow.python.data.kernel_tests import test_base
|
|
from tensorflow.python.data.ops import dataset_ops
|
|
from tensorflow.python.data.util import nest
|
|
from tensorflow.python.data.util import structure
|
|
from tensorflow.python.framework import combinations
|
|
from tensorflow.python.framework import constant_op
|
|
from tensorflow.python.framework import dtypes
|
|
from tensorflow.python.framework import sparse_tensor
|
|
from tensorflow.python.framework import tensor
|
|
from tensorflow.python.framework import tensor_shape
|
|
from tensorflow.python.ops import array_ops
|
|
from tensorflow.python.ops import tensor_array_ops
|
|
from tensorflow.python.ops import variables
|
|
from tensorflow.python.ops.ragged import ragged_factory_ops
|
|
from tensorflow.python.ops.ragged import ragged_tensor
|
|
from tensorflow.python.ops.ragged import ragged_tensor_value
|
|
from tensorflow.python.platform import test
|
|
from tensorflow.python.util.compat import collections_abc
|
|
|
|
# NOTE(mrry): Arguments of parameterized tests are lifted into lambdas to make
|
|
# sure they are not executed before the (eager- or graph-mode) test environment
|
|
# has been set up.
|
|
#
|
|
|
|
|
|
def _test_flat_structure_combinations():
|
|
cases = [
|
|
("Tensor", lambda: constant_op.constant(37.0),
|
|
lambda: tensor.TensorSpec, lambda: [dtypes.float32], lambda: [[]]),
|
|
("TensorArray", lambda: tensor_array_ops.TensorArray(
|
|
dtype=dtypes.float32, element_shape=(3,), size=0),
|
|
lambda: tensor_array_ops.TensorArraySpec, lambda: [dtypes.variant],
|
|
lambda: [[]]),
|
|
("SparseTensor", lambda: sparse_tensor.SparseTensor(
|
|
indices=[[3, 4]], values=[-1], dense_shape=[4, 5]),
|
|
lambda: sparse_tensor.SparseTensorSpec, lambda: [dtypes.variant],
|
|
lambda: [None]),
|
|
("RaggedTensor", lambda: ragged_factory_ops.constant([[1, 2], [], [4]]),
|
|
lambda: ragged_tensor.RaggedTensorSpec, lambda: [dtypes.variant],
|
|
lambda: [None]),
|
|
("Nested_0", lambda:
|
|
(constant_op.constant(37.0), constant_op.constant([1, 2, 3])),
|
|
lambda: tuple, lambda: [dtypes.float32, dtypes.int32],
|
|
lambda: [[], [3]]),
|
|
("Nested_1", lambda: {
|
|
"a": constant_op.constant(37.0),
|
|
"b": constant_op.constant([1, 2, 3])
|
|
}, lambda: dict, lambda: [dtypes.float32, dtypes.int32],
|
|
lambda: [[], [3]]),
|
|
("Nested_2", lambda: {
|
|
"a":
|
|
constant_op.constant(37.0),
|
|
"b": (sparse_tensor.SparseTensor(
|
|
indices=[[0, 0]], values=[1], dense_shape=[1, 1]),
|
|
sparse_tensor.SparseTensor(
|
|
indices=[[3, 4]], values=[-1], dense_shape=[4, 5]))
|
|
}, lambda: dict, lambda: [dtypes.float32, dtypes.variant, dtypes.variant],
|
|
lambda: [[], None, None]),
|
|
]
|
|
|
|
def reduce_fn(x, y):
|
|
# workaround for long line
|
|
name, value_fn = y[:2]
|
|
expected_structure_fn, expected_types_fn, expected_shapes_fn = y[2:]
|
|
return x + combinations.combine(
|
|
value_fn=combinations.NamedObject("value_fn.{}".format(name), value_fn),
|
|
expected_structure_fn=combinations.NamedObject(
|
|
"expected_structure_fn.{}".format(name), expected_structure_fn),
|
|
expected_types_fn=combinations.NamedObject(
|
|
"expected_types_fn.{}".format(name), expected_types_fn),
|
|
expected_shapes_fn=combinations.NamedObject(
|
|
"expected_shapes_fn.{}".format(name), expected_shapes_fn))
|
|
|
|
return functools.reduce(reduce_fn, cases, [])
|
|
|
|
|
|
def _test_is_compatible_with_structure_combinations():
|
|
cases = [
|
|
("Tensor", lambda: constant_op.constant(37.0), lambda: [
|
|
constant_op.constant(38.0),
|
|
array_ops.placeholder(dtypes.float32), 42.0,
|
|
np.array(42.0, dtype=np.float32)
|
|
], lambda: [constant_op.constant([1.0, 2.0]),
|
|
constant_op.constant(37)]),
|
|
# TODO(b/209081027): add Python constant and TF constant to the
|
|
# incompatible branch after ResourceVariable becoming a CompositeTensor.
|
|
("Variable", lambda: variables.Variable(100.0),
|
|
lambda: [variables.Variable(99.0)],
|
|
lambda: [1]),
|
|
("TensorArray", lambda: tensor_array_ops.TensorArray(
|
|
dtype=dtypes.float32, element_shape=(3,), size=0), lambda: [
|
|
tensor_array_ops.TensorArray(
|
|
dtype=dtypes.float32, element_shape=(3,), size=0),
|
|
tensor_array_ops.TensorArray(
|
|
dtype=dtypes.float32, element_shape=(3,), size=10)
|
|
], lambda: [
|
|
tensor_array_ops.TensorArray(
|
|
dtype=dtypes.int32, element_shape=(3,), size=0),
|
|
tensor_array_ops.TensorArray(
|
|
dtype=dtypes.float32, element_shape=(), size=0)
|
|
]),
|
|
("SparseTensor", lambda: sparse_tensor.SparseTensor(
|
|
indices=[[3, 4]], values=[-1], dense_shape=[4, 5]),
|
|
lambda: [
|
|
sparse_tensor.SparseTensor(
|
|
indices=[[1, 1], [3, 4]], values=[10, -1], dense_shape=[4, 5]),
|
|
sparse_tensor.SparseTensorValue(
|
|
indices=[[1, 1], [3, 4]], values=[10, -1], dense_shape=[4, 5]),
|
|
array_ops.sparse_placeholder(dtype=dtypes.int32),
|
|
array_ops.sparse_placeholder(dtype=dtypes.int32, shape=[None, None])
|
|
], lambda: [
|
|
constant_op.constant(37, shape=[4, 5]),
|
|
sparse_tensor.SparseTensor(
|
|
indices=[[3, 4]], values=[-1], dense_shape=[5, 6]),
|
|
array_ops.sparse_placeholder(
|
|
dtype=dtypes.int32, shape=[None, None, None]),
|
|
sparse_tensor.SparseTensor(
|
|
indices=[[3, 4]], values=[-1.0], dense_shape=[4, 5])
|
|
]),
|
|
("RaggedTensor", lambda: ragged_factory_ops.constant([[1, 2], [], [3]]),
|
|
lambda: [
|
|
ragged_factory_ops.constant([[1, 2], [3, 4], []]),
|
|
ragged_factory_ops.constant([[1], [2, 3, 4], [5]]),
|
|
], lambda: [
|
|
ragged_factory_ops.constant(1),
|
|
ragged_factory_ops.constant([1, 2]),
|
|
ragged_factory_ops.constant([[1], [2]]),
|
|
ragged_factory_ops.constant([["a", "b"]]),
|
|
]),
|
|
("Nested", lambda: {
|
|
"a": constant_op.constant(37.0),
|
|
"b": constant_op.constant([1, 2, 3])
|
|
}, lambda: [{
|
|
"a": constant_op.constant(15.0),
|
|
"b": constant_op.constant([4, 5, 6])
|
|
}], lambda: [{
|
|
"a": constant_op.constant(15.0),
|
|
"b": constant_op.constant([4, 5, 6, 7])
|
|
}, {
|
|
"a": constant_op.constant(15),
|
|
"b": constant_op.constant([4, 5, 6])
|
|
}, {
|
|
"a":
|
|
constant_op.constant(15),
|
|
"b":
|
|
sparse_tensor.SparseTensor(
|
|
indices=[[0], [1], [2]], values=[4, 5, 6], dense_shape=[3])
|
|
}, (constant_op.constant(15.0), constant_op.constant([4, 5, 6]))]),
|
|
]
|
|
|
|
def reduce_fn(x, y):
|
|
name, original_value_fn, compatible_values_fn, incompatible_values_fn = y
|
|
return x + combinations.combine(
|
|
original_value_fn=combinations.NamedObject(
|
|
"original_value_fn.{}".format(name), original_value_fn),
|
|
compatible_values_fn=combinations.NamedObject(
|
|
"compatible_values_fn.{}".format(name), compatible_values_fn),
|
|
incompatible_values_fn=combinations.NamedObject(
|
|
"incompatible_values_fn.{}".format(name), incompatible_values_fn))
|
|
|
|
return functools.reduce(reduce_fn, cases, [])
|
|
|
|
|
|
def _test_structure_from_value_equality_combinations():
|
|
cases = [
|
|
("Tensor", lambda: constant_op.constant(37.0),
|
|
lambda: constant_op.constant(42.0), lambda: constant_op.constant([5])),
|
|
("TensorArray", lambda: tensor_array_ops.TensorArray(
|
|
dtype=dtypes.float32, element_shape=(3,), size=0),
|
|
lambda: tensor_array_ops.TensorArray(
|
|
dtype=dtypes.float32, element_shape=(3,), size=0),
|
|
lambda: tensor_array_ops.TensorArray(
|
|
dtype=dtypes.int32, element_shape=(), size=0)),
|
|
("SparseTensor", lambda: sparse_tensor.SparseTensor(
|
|
indices=[[3, 4]], values=[-1], dense_shape=[4, 5]),
|
|
lambda: sparse_tensor.SparseTensor(
|
|
indices=[[1, 2]], values=[42], dense_shape=[4, 5]), lambda:
|
|
sparse_tensor.SparseTensor(indices=[[3]], values=[-1], dense_shape=[5]),
|
|
lambda: sparse_tensor.SparseTensor(
|
|
indices=[[3, 4]], values=[1.0], dense_shape=[4, 5])),
|
|
("RaggedTensor", lambda: ragged_factory_ops.constant([[[1, 2]], [[3]]]),
|
|
lambda: ragged_factory_ops.constant([[[5]], [[8], [3, 2]]]),
|
|
lambda: ragged_factory_ops.constant([[[1]], [[2], [3]]], ragged_rank=1),
|
|
lambda: ragged_factory_ops.constant([[[1.0, 2.0]], [[3.0]]]),
|
|
lambda: ragged_factory_ops.constant([[[1]], [[2]], [[3]]])),
|
|
("Nested", lambda: {
|
|
"a": constant_op.constant(37.0),
|
|
"b": constant_op.constant([1, 2, 3])
|
|
}, lambda: {
|
|
"a": constant_op.constant(42.0),
|
|
"b": constant_op.constant([4, 5, 6])
|
|
}, lambda: {
|
|
"a": constant_op.constant([1, 2, 3]),
|
|
"b": constant_op.constant(37.0)
|
|
}),
|
|
]
|
|
|
|
def reduce_fn(x, y):
|
|
name, value1_fn, value2_fn, *not_equal_value_fns = y
|
|
return x + combinations.combine(
|
|
value1_fn=combinations.NamedObject("value1_fn.{}".format(name),
|
|
value1_fn),
|
|
value2_fn=combinations.NamedObject("value2_fn.{}".format(name),
|
|
value2_fn),
|
|
not_equal_value_fns=combinations.NamedObject(
|
|
"not_equal_value_fns.{}".format(name), not_equal_value_fns))
|
|
|
|
return functools.reduce(reduce_fn, cases, [])
|
|
|
|
|
|
def _test_ragged_structure_inequality_combinations():
|
|
cases = [
|
|
(ragged_tensor.RaggedTensorSpec(None, dtypes.int32, 1),
|
|
ragged_tensor.RaggedTensorSpec(None, dtypes.int32, 2)),
|
|
(ragged_tensor.RaggedTensorSpec([3, None], dtypes.int32, 1),
|
|
ragged_tensor.RaggedTensorSpec([5, None], dtypes.int32, 1)),
|
|
(ragged_tensor.RaggedTensorSpec(None, dtypes.int32, 1),
|
|
ragged_tensor.RaggedTensorSpec(None, dtypes.float32, 1)),
|
|
]
|
|
|
|
def reduce_fn(x, y):
|
|
spec1, spec2 = y
|
|
return x + combinations.combine(spec1=spec1, spec2=spec2)
|
|
|
|
return functools.reduce(reduce_fn, cases, [])
|
|
|
|
|
|
def _test_hash_combinations():
|
|
cases = [
|
|
("Tensor", lambda: constant_op.constant(37.0),
|
|
lambda: constant_op.constant(42.0), lambda: constant_op.constant([5])),
|
|
("TensorArray", lambda: tensor_array_ops.TensorArray(
|
|
dtype=dtypes.float32, element_shape=(3,), size=0),
|
|
lambda: tensor_array_ops.TensorArray(
|
|
dtype=dtypes.float32, element_shape=(3,), size=0),
|
|
lambda: tensor_array_ops.TensorArray(
|
|
dtype=dtypes.int32, element_shape=(), size=0)),
|
|
("SparseTensor", lambda: sparse_tensor.SparseTensor(
|
|
indices=[[3, 4]], values=[-1], dense_shape=[4, 5]),
|
|
lambda: sparse_tensor.SparseTensor(
|
|
indices=[[1, 2]], values=[42], dense_shape=[4, 5]), lambda:
|
|
sparse_tensor.SparseTensor(indices=[[3]], values=[-1], dense_shape=[5])),
|
|
("Nested", lambda: {
|
|
"a": constant_op.constant(37.0),
|
|
"b": constant_op.constant([1, 2, 3])
|
|
}, lambda: {
|
|
"a": constant_op.constant(42.0),
|
|
"b": constant_op.constant([4, 5, 6])
|
|
}, lambda: {
|
|
"a": constant_op.constant([1, 2, 3]),
|
|
"b": constant_op.constant(37.0)
|
|
}),
|
|
]
|
|
|
|
def reduce_fn(x, y):
|
|
name, value1_fn, value2_fn, value3_fn = y
|
|
return x + combinations.combine(
|
|
value1_fn=combinations.NamedObject("value1_fn.{}".format(name),
|
|
value1_fn),
|
|
value2_fn=combinations.NamedObject("value2_fn.{}".format(name),
|
|
value2_fn),
|
|
value3_fn=combinations.NamedObject("value3_fn.{}".format(name),
|
|
value3_fn))
|
|
|
|
return functools.reduce(reduce_fn, cases, [])
|
|
|
|
|
|
def _test_round_trip_conversion_combinations():
|
|
cases = [
|
|
(
|
|
"Tensor",
|
|
lambda: constant_op.constant(37.0),
|
|
),
|
|
(
|
|
"SparseTensor",
|
|
lambda: sparse_tensor.SparseTensor(
|
|
indices=[[3, 4]], values=[-1], dense_shape=[4, 5]),
|
|
),
|
|
("TensorArray", lambda: tensor_array_ops.TensorArray(
|
|
dtype=dtypes.float32, element_shape=(), size=1).write(0, 7)),
|
|
(
|
|
"RaggedTensor",
|
|
lambda: ragged_factory_ops.constant([[1, 2], [], [3]]),
|
|
),
|
|
(
|
|
"Nested_0",
|
|
lambda: {
|
|
"a": constant_op.constant(37.0),
|
|
"b": constant_op.constant([1, 2, 3])
|
|
},
|
|
),
|
|
(
|
|
"Nested_1",
|
|
lambda: {
|
|
"a":
|
|
constant_op.constant(37.0),
|
|
"b": (sparse_tensor.SparseTensor(
|
|
indices=[[0, 0]], values=[1], dense_shape=[1, 1]),
|
|
sparse_tensor.SparseTensor(
|
|
indices=[[3, 4]], values=[-1], dense_shape=[4, 5]))
|
|
},
|
|
),
|
|
]
|
|
|
|
def reduce_fn(x, y):
|
|
name, value_fn = y
|
|
return x + combinations.combine(
|
|
value_fn=combinations.NamedObject("value_fn.{}".format(name), value_fn))
|
|
|
|
return functools.reduce(reduce_fn, cases, [])
|
|
|
|
|
|
def _test_convert_legacy_structure_combinations():
|
|
cases = [
|
|
(dtypes.float32, tensor_shape.TensorShape([]), tensor.Tensor,
|
|
tensor.TensorSpec([], dtypes.float32)),
|
|
(dtypes.int32, tensor_shape.TensorShape([2,
|
|
2]), sparse_tensor.SparseTensor,
|
|
sparse_tensor.SparseTensorSpec([2, 2], dtypes.int32)),
|
|
(dtypes.int32, tensor_shape.TensorShape([None, True, 2, 2]),
|
|
tensor_array_ops.TensorArray,
|
|
tensor_array_ops.TensorArraySpec([2, 2],
|
|
dtypes.int32,
|
|
dynamic_size=None,
|
|
infer_shape=True)),
|
|
(dtypes.int32, tensor_shape.TensorShape([True, None, 2, 2]),
|
|
tensor_array_ops.TensorArray,
|
|
tensor_array_ops.TensorArraySpec([2, 2],
|
|
dtypes.int32,
|
|
dynamic_size=True,
|
|
infer_shape=None)),
|
|
(dtypes.int32, tensor_shape.TensorShape([True, False, 2, 2]),
|
|
tensor_array_ops.TensorArray,
|
|
tensor_array_ops.TensorArraySpec([2, 2],
|
|
dtypes.int32,
|
|
dynamic_size=True,
|
|
infer_shape=False)),
|
|
(dtypes.int32, tensor_shape.TensorShape([2, None]),
|
|
ragged_tensor.RaggedTensorSpec([2, None], dtypes.int32, 1),
|
|
ragged_tensor.RaggedTensorSpec([2, None], dtypes.int32, 1)),
|
|
({
|
|
"a": dtypes.float32,
|
|
"b": (dtypes.int32, dtypes.string)
|
|
}, {
|
|
"a": tensor_shape.TensorShape([]),
|
|
"b": (tensor_shape.TensorShape([2, 2]), tensor_shape.TensorShape([]))
|
|
}, {
|
|
"a": tensor.Tensor,
|
|
"b": (sparse_tensor.SparseTensor, tensor.Tensor)
|
|
}, {
|
|
"a":
|
|
tensor.TensorSpec([], dtypes.float32),
|
|
"b": (sparse_tensor.SparseTensorSpec([2, 2], dtypes.int32),
|
|
tensor.TensorSpec([], dtypes.string))
|
|
})
|
|
]
|
|
|
|
def reduce_fn(x, y):
|
|
output_types, output_shapes, output_classes, expected_structure = y
|
|
return x + combinations.combine(
|
|
output_types=output_types,
|
|
output_shapes=output_shapes,
|
|
output_classes=output_classes,
|
|
expected_structure=expected_structure)
|
|
|
|
return functools.reduce(reduce_fn, cases, [])
|
|
|
|
|
|
def _test_batch_combinations():
|
|
cases = [
|
|
(tensor.TensorSpec([], dtypes.float32), 32,
|
|
tensor.TensorSpec([32], dtypes.float32)),
|
|
(tensor.TensorSpec([], dtypes.float32), None,
|
|
tensor.TensorSpec([None], dtypes.float32)),
|
|
(sparse_tensor.SparseTensorSpec([None], dtypes.float32), 32,
|
|
sparse_tensor.SparseTensorSpec([32, None], dtypes.float32)),
|
|
(sparse_tensor.SparseTensorSpec([4], dtypes.float32), None,
|
|
sparse_tensor.SparseTensorSpec([None, 4], dtypes.float32)),
|
|
(ragged_tensor.RaggedTensorSpec([2, None], dtypes.float32, 1), 32,
|
|
ragged_tensor.RaggedTensorSpec([32, 2, None], dtypes.float32, 2)),
|
|
(ragged_tensor.RaggedTensorSpec([4, None], dtypes.float32, 1), None,
|
|
ragged_tensor.RaggedTensorSpec([None, 4, None], dtypes.float32, 2)),
|
|
({
|
|
"a":
|
|
tensor.TensorSpec([], dtypes.float32),
|
|
"b": (sparse_tensor.SparseTensorSpec([2, 2], dtypes.int32),
|
|
tensor.TensorSpec([], dtypes.string))
|
|
}, 128, {
|
|
"a":
|
|
tensor.TensorSpec([128], dtypes.float32),
|
|
"b": (sparse_tensor.SparseTensorSpec([128, 2, 2], dtypes.int32),
|
|
tensor.TensorSpec([128], dtypes.string))
|
|
}),
|
|
]
|
|
|
|
def reduce_fn(x, y):
|
|
element_structure, batch_size, expected_batched_structure = y
|
|
return x + combinations.combine(
|
|
element_structure=element_structure,
|
|
batch_size=batch_size,
|
|
expected_batched_structure=expected_batched_structure)
|
|
|
|
return functools.reduce(reduce_fn, cases, [])
|
|
|
|
|
|
def _test_unbatch_combinations():
|
|
cases = [
|
|
(tensor.TensorSpec([32], dtypes.float32),
|
|
tensor.TensorSpec([], dtypes.float32)),
|
|
(tensor.TensorSpec([None], dtypes.float32),
|
|
tensor.TensorSpec([], dtypes.float32)),
|
|
(sparse_tensor.SparseTensorSpec([32, None], dtypes.float32),
|
|
sparse_tensor.SparseTensorSpec([None], dtypes.float32)),
|
|
(sparse_tensor.SparseTensorSpec([None, 4], dtypes.float32),
|
|
sparse_tensor.SparseTensorSpec([4], dtypes.float32)),
|
|
(ragged_tensor.RaggedTensorSpec([32, None, None], dtypes.float32, 2),
|
|
ragged_tensor.RaggedTensorSpec([None, None], dtypes.float32, 1)),
|
|
(ragged_tensor.RaggedTensorSpec([None, None, None], dtypes.float32, 2),
|
|
ragged_tensor.RaggedTensorSpec([None, None], dtypes.float32, 1)),
|
|
({
|
|
"a":
|
|
tensor.TensorSpec([128], dtypes.float32),
|
|
"b": (sparse_tensor.SparseTensorSpec([128, 2, 2], dtypes.int32),
|
|
tensor.TensorSpec([None], dtypes.string))
|
|
}, {
|
|
"a":
|
|
tensor.TensorSpec([], dtypes.float32),
|
|
"b": (sparse_tensor.SparseTensorSpec([2, 2], dtypes.int32),
|
|
tensor.TensorSpec([], dtypes.string))
|
|
}),
|
|
]
|
|
|
|
def reduce_fn(x, y):
|
|
element_structure, expected_unbatched_structure = y
|
|
return x + combinations.combine(
|
|
element_structure=element_structure,
|
|
expected_unbatched_structure=expected_unbatched_structure)
|
|
|
|
return functools.reduce(reduce_fn, cases, [])
|
|
|
|
|
|
def _test_to_batched_tensor_list_combinations():
|
|
cases = [
|
|
("Tensor", lambda: constant_op.constant([[1.0, 2.0], [3.0, 4.0]]),
|
|
lambda: constant_op.constant([1.0, 2.0])),
|
|
("SparseTensor", lambda: sparse_tensor.SparseTensor(
|
|
indices=[[0, 0], [1, 1]], values=[13, 27], dense_shape=[2, 2]),
|
|
lambda: sparse_tensor.SparseTensor(
|
|
indices=[[0]], values=[13], dense_shape=[2])),
|
|
("RaggedTensor", lambda: ragged_factory_ops.constant([[[1]], [[2]]]),
|
|
lambda: ragged_factory_ops.constant([[1]])),
|
|
("Nest", lambda:
|
|
(constant_op.constant([[1.0, 2.0], [3.0, 4.0]]),
|
|
sparse_tensor.SparseTensor(
|
|
indices=[[0, 0], [1, 1]], values=[13, 27], dense_shape=[2, 2])),
|
|
lambda:
|
|
(constant_op.constant([1.0, 2.0]),
|
|
sparse_tensor.SparseTensor(indices=[[0]], values=[13], dense_shape=[2]))
|
|
),
|
|
]
|
|
|
|
def reduce_fn(x, y):
|
|
name, value_fn, element_0_fn = y
|
|
return x + combinations.combine(
|
|
value_fn=combinations.NamedObject("value_fn.{}".format(name), value_fn),
|
|
element_0_fn=combinations.NamedObject("element_0_fn.{}".format(name),
|
|
element_0_fn))
|
|
|
|
return functools.reduce(reduce_fn, cases, [])
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class MaskedTensor:
|
|
mask: bool
|
|
value: tensor.Tensor
|
|
|
|
def __tf_flatten__(self):
|
|
metadata = (self.mask,)
|
|
components = (self.value,)
|
|
return metadata, components
|
|
|
|
def __tf_unflatten__(self, metadata, components):
|
|
mask = metadata[0]
|
|
value = components[0]
|
|
return MaskedTensor(mask=mask, value=value)
|
|
|
|
|
|
# TODO(jsimsa): Add tests for OptionalStructure and DatasetStructure.
|
|
class StructureTest(test_base.DatasetTestBase, parameterized.TestCase):
|
|
|
|
# pylint: disable=g-long-lambda,protected-access
|
|
@combinations.generate(
|
|
combinations.times(test_base.default_test_combinations(),
|
|
_test_flat_structure_combinations()))
|
|
def testFlatStructure(self, value_fn, expected_structure_fn,
|
|
expected_types_fn, expected_shapes_fn):
|
|
value = value_fn()
|
|
expected_structure = expected_structure_fn()
|
|
expected_types = expected_types_fn()
|
|
expected_shapes = expected_shapes_fn()
|
|
s = structure.type_spec_from_value(value)
|
|
self.assertIsInstance(s, expected_structure)
|
|
flat_types = structure.get_flat_tensor_types(s)
|
|
self.assertEqual(expected_types, flat_types)
|
|
flat_shapes = structure.get_flat_tensor_shapes(s)
|
|
self.assertLen(flat_shapes, len(expected_shapes))
|
|
for expected, actual in zip(expected_shapes, flat_shapes):
|
|
if expected is None:
|
|
self.assertEqual(actual.ndims, None)
|
|
else:
|
|
self.assertEqual(actual.as_list(), expected)
|
|
|
|
@combinations.generate(
|
|
combinations.times(test_base.graph_only_combinations(),
|
|
_test_is_compatible_with_structure_combinations()))
|
|
def testIsCompatibleWithStructure(self, original_value_fn,
|
|
compatible_values_fn,
|
|
incompatible_values_fn):
|
|
original_value = original_value_fn()
|
|
compatible_values = compatible_values_fn()
|
|
incompatible_values = incompatible_values_fn()
|
|
|
|
s = structure.type_spec_from_value(original_value)
|
|
for compatible_value in compatible_values:
|
|
self.assertTrue(
|
|
structure.are_compatible(
|
|
s, structure.type_spec_from_value(compatible_value)))
|
|
for incompatible_value in incompatible_values:
|
|
self.assertFalse(
|
|
structure.are_compatible(
|
|
s, structure.type_spec_from_value(incompatible_value)))
|
|
|
|
@combinations.generate(
|
|
combinations.times(test_base.default_test_combinations(),
|
|
_test_structure_from_value_equality_combinations()))
|
|
def testStructureFromValueEquality(self, value1_fn, value2_fn,
|
|
not_equal_value_fns):
|
|
# pylint: disable=g-generic-assert
|
|
not_equal_value_fns = not_equal_value_fns._obj
|
|
s1 = structure.type_spec_from_value(value1_fn())
|
|
s2 = structure.type_spec_from_value(value2_fn())
|
|
self.assertEqual(s1, s1) # check __eq__ operator.
|
|
self.assertEqual(s1, s2) # check __eq__ operator.
|
|
self.assertFalse(s1 != s1) # check __ne__ operator.
|
|
self.assertFalse(s1 != s2) # check __ne__ operator.
|
|
for c1, c2 in zip(nest.flatten(s1), nest.flatten(s2)):
|
|
self.assertEqual(hash(c1), hash(c1))
|
|
self.assertEqual(hash(c1), hash(c2))
|
|
for value_fn in not_equal_value_fns:
|
|
s3 = structure.type_spec_from_value(value_fn())
|
|
self.assertNotEqual(s1, s3) # check __ne__ operator.
|
|
self.assertNotEqual(s2, s3) # check __ne__ operator.
|
|
self.assertFalse(s1 == s3) # check __eq_ operator.
|
|
self.assertFalse(s2 == s3) # check __eq_ operator.
|
|
|
|
@combinations.generate(
|
|
combinations.times(test_base.default_test_combinations(),
|
|
_test_ragged_structure_inequality_combinations()))
|
|
def testRaggedStructureInequality(self, spec1, spec2):
|
|
# pylint: disable=g-generic-assert
|
|
self.assertNotEqual(spec1, spec2) # check __ne__ operator.
|
|
self.assertFalse(spec1 == spec2) # check __eq__ operator.
|
|
|
|
@combinations.generate(
|
|
combinations.times(test_base.default_test_combinations(),
|
|
_test_hash_combinations()))
|
|
def testHash(self, value1_fn, value2_fn, value3_fn):
|
|
s1 = structure.type_spec_from_value(value1_fn())
|
|
s2 = structure.type_spec_from_value(value2_fn())
|
|
s3 = structure.type_spec_from_value(value3_fn())
|
|
for c1, c2, c3 in zip(nest.flatten(s1), nest.flatten(s2), nest.flatten(s3)):
|
|
self.assertEqual(hash(c1), hash(c1))
|
|
self.assertEqual(hash(c1), hash(c2))
|
|
self.assertNotEqual(hash(c1), hash(c3))
|
|
self.assertNotEqual(hash(c2), hash(c3))
|
|
|
|
@combinations.generate(
|
|
combinations.times(test_base.default_test_combinations(),
|
|
_test_round_trip_conversion_combinations()))
|
|
def testRoundTripConversion(self, value_fn):
|
|
value = value_fn()
|
|
s = structure.type_spec_from_value(value)
|
|
|
|
def maybe_stack_ta(v):
|
|
if isinstance(v, tensor_array_ops.TensorArray):
|
|
return v.stack()
|
|
return v
|
|
|
|
before = self.evaluate(maybe_stack_ta(value))
|
|
after = self.evaluate(
|
|
maybe_stack_ta(
|
|
structure.from_tensor_list(s, structure.to_tensor_list(s, value))))
|
|
|
|
flat_before = nest.flatten(before)
|
|
flat_after = nest.flatten(after)
|
|
for b, a in zip(flat_before, flat_after):
|
|
if isinstance(b, sparse_tensor.SparseTensorValue):
|
|
self.assertAllEqual(b.indices, a.indices)
|
|
self.assertAllEqual(b.values, a.values)
|
|
self.assertAllEqual(b.dense_shape, a.dense_shape)
|
|
elif isinstance(
|
|
b,
|
|
(ragged_tensor.RaggedTensor, ragged_tensor_value.RaggedTensorValue)):
|
|
self.assertAllEqual(b, a)
|
|
else:
|
|
self.assertAllEqual(b, a)
|
|
|
|
# pylint: enable=g-long-lambda
|
|
|
|
def preserveStaticShape(self):
|
|
rt = ragged_factory_ops.constant([[1, 2], [], [3]])
|
|
rt_s = structure.type_spec_from_value(rt)
|
|
rt_after = structure.from_tensor_list(rt_s,
|
|
structure.to_tensor_list(rt_s, rt))
|
|
self.assertEqual(rt_after.row_splits.shape.as_list(),
|
|
rt.row_splits.shape.as_list())
|
|
self.assertEqual(rt_after.values.shape.as_list(), [None])
|
|
|
|
st = sparse_tensor.SparseTensor(
|
|
indices=[[3, 4]], values=[-1], dense_shape=[4, 5])
|
|
st_s = structure.type_spec_from_value(st)
|
|
st_after = structure.from_tensor_list(st_s,
|
|
structure.to_tensor_list(st_s, st))
|
|
self.assertEqual(st_after.indices.shape.as_list(), [None, 2])
|
|
self.assertEqual(st_after.values.shape.as_list(), [None])
|
|
self.assertEqual(st_after.dense_shape.shape.as_list(),
|
|
st.dense_shape.shape.as_list())
|
|
|
|
@combinations.generate(test_base.default_test_combinations())
|
|
def testPreserveTensorArrayShape(self):
|
|
ta = tensor_array_ops.TensorArray(
|
|
dtype=dtypes.int32, size=1, element_shape=(3,))
|
|
ta_s = structure.type_spec_from_value(ta)
|
|
ta_after = structure.from_tensor_list(ta_s,
|
|
structure.to_tensor_list(ta_s, ta))
|
|
self.assertEqual(ta_after.element_shape.as_list(), [3])
|
|
|
|
@combinations.generate(test_base.default_test_combinations())
|
|
def testPreserveInferredTensorArrayShape(self):
|
|
ta = tensor_array_ops.TensorArray(dtype=dtypes.int32, size=1)
|
|
# Shape is inferred from the write.
|
|
ta = ta.write(0, [1, 2, 3])
|
|
ta_s = structure.type_spec_from_value(ta)
|
|
ta_after = structure.from_tensor_list(ta_s,
|
|
structure.to_tensor_list(ta_s, ta))
|
|
self.assertEqual(ta_after.element_shape.as_list(), [3])
|
|
|
|
@combinations.generate(test_base.default_test_combinations())
|
|
def testIncompatibleStructure(self):
|
|
# Define three mutually incompatible values/structures, and assert that:
|
|
# 1. Using one structure to flatten a value with an incompatible structure
|
|
# fails.
|
|
# 2. Using one structure to restructure a flattened value with an
|
|
# incompatible structure fails.
|
|
value_tensor = constant_op.constant(42.0)
|
|
s_tensor = structure.type_spec_from_value(value_tensor)
|
|
flat_tensor = structure.to_tensor_list(s_tensor, value_tensor)
|
|
|
|
value_sparse_tensor = sparse_tensor.SparseTensor(
|
|
indices=[[0, 0]], values=[1], dense_shape=[1, 1])
|
|
s_sparse_tensor = structure.type_spec_from_value(value_sparse_tensor)
|
|
flat_sparse_tensor = structure.to_tensor_list(s_sparse_tensor,
|
|
value_sparse_tensor)
|
|
|
|
value_nest = {
|
|
"a": constant_op.constant(37.0),
|
|
"b": constant_op.constant([1, 2, 3])
|
|
}
|
|
s_nest = structure.type_spec_from_value(value_nest)
|
|
flat_nest = structure.to_tensor_list(s_nest, value_nest)
|
|
|
|
with self.assertRaisesRegex(
|
|
ValueError, r"SparseTensor.* is not convertible to a tensor with "
|
|
r"dtype.*float32.* and shape \(\)"):
|
|
structure.to_tensor_list(s_tensor, value_sparse_tensor)
|
|
with self.assertRaisesRegex(
|
|
ValueError, "The two structures don't have the same nested structure."):
|
|
structure.to_tensor_list(s_tensor, value_nest)
|
|
|
|
with self.assertRaisesRegex(TypeError,
|
|
"neither a SparseTensor nor SparseTensorValue"):
|
|
structure.to_tensor_list(s_sparse_tensor, value_tensor)
|
|
|
|
with self.assertRaisesRegex(
|
|
ValueError, "The two structures don't have the same nested structure."):
|
|
structure.to_tensor_list(s_sparse_tensor, value_nest)
|
|
|
|
with self.assertRaisesRegex(
|
|
ValueError, "The two structures don't have the same nested structure."):
|
|
structure.to_tensor_list(s_nest, value_tensor)
|
|
|
|
with self.assertRaisesRegex(
|
|
ValueError, "The two structures don't have the same nested structure."):
|
|
structure.to_tensor_list(s_nest, value_sparse_tensor)
|
|
|
|
with self.assertRaisesRegex(
|
|
ValueError,
|
|
"Cannot create a Tensor from the tensor list because item 0 "
|
|
".*tf.Tensor.* is incompatible with the expected TypeSpec "
|
|
".*TensorSpec.*"):
|
|
structure.from_tensor_list(s_tensor, flat_sparse_tensor)
|
|
|
|
with self.assertRaisesRegex(ValueError, "Expected 1 tensors but got 2."):
|
|
structure.from_tensor_list(s_tensor, flat_nest)
|
|
|
|
with self.assertRaisesRegex(
|
|
ValueError, "Cannot create a SparseTensor from the tensor list because "
|
|
"item 0 .*tf.Tensor.* is incompatible with the expected TypeSpec "
|
|
".*TensorSpec.*"):
|
|
structure.from_tensor_list(s_sparse_tensor, flat_tensor)
|
|
|
|
with self.assertRaisesRegex(ValueError, "Expected 1 tensors but got 2."):
|
|
structure.from_tensor_list(s_sparse_tensor, flat_nest)
|
|
|
|
with self.assertRaisesRegex(ValueError, "Expected 2 tensors but got 1."):
|
|
structure.from_tensor_list(s_nest, flat_tensor)
|
|
|
|
with self.assertRaisesRegex(ValueError, "Expected 2 tensors but got 1."):
|
|
structure.from_tensor_list(s_nest, flat_sparse_tensor)
|
|
|
|
@combinations.generate(test_base.default_test_combinations())
|
|
def testIncompatibleNestedStructure(self):
|
|
# Define three mutually incompatible nested values/structures, and assert
|
|
# that:
|
|
# 1. Using one structure to flatten a value with an incompatible structure
|
|
# fails.
|
|
# 2. Using one structure to restructure a flattened value with an
|
|
# incompatible structure fails.
|
|
|
|
value_0 = {
|
|
"a": constant_op.constant(37.0),
|
|
"b": constant_op.constant([1, 2, 3])
|
|
}
|
|
s_0 = structure.type_spec_from_value(value_0)
|
|
flat_s_0 = structure.to_tensor_list(s_0, value_0)
|
|
|
|
# `value_1` has compatible nested structure with `value_0`, but different
|
|
# classes.
|
|
value_1 = {
|
|
"a":
|
|
constant_op.constant(37.0),
|
|
"b":
|
|
sparse_tensor.SparseTensor(
|
|
indices=[[0, 0]], values=[1], dense_shape=[1, 1])
|
|
}
|
|
s_1 = structure.type_spec_from_value(value_1)
|
|
flat_s_1 = structure.to_tensor_list(s_1, value_1)
|
|
|
|
# `value_2` has incompatible nested structure with `value_0` and `value_1`.
|
|
value_2 = {
|
|
"a":
|
|
constant_op.constant(37.0),
|
|
"b": (sparse_tensor.SparseTensor(
|
|
indices=[[0, 0]], values=[1], dense_shape=[1, 1]),
|
|
sparse_tensor.SparseTensor(
|
|
indices=[[3, 4]], values=[-1], dense_shape=[4, 5]))
|
|
}
|
|
s_2 = structure.type_spec_from_value(value_2)
|
|
flat_s_2 = structure.to_tensor_list(s_2, value_2)
|
|
|
|
with self.assertRaisesRegex(
|
|
ValueError, r"SparseTensor.* is not convertible to a tensor with "
|
|
r"dtype.*int32.* and shape \(3,\)"):
|
|
structure.to_tensor_list(s_0, value_1)
|
|
|
|
with self.assertRaisesRegex(
|
|
ValueError, "The two structures don't have the same nested structure."):
|
|
structure.to_tensor_list(s_0, value_2)
|
|
|
|
with self.assertRaisesRegex(TypeError,
|
|
"neither a SparseTensor nor SparseTensorValue"):
|
|
structure.to_tensor_list(s_1, value_0)
|
|
|
|
with self.assertRaisesRegex(
|
|
ValueError, "The two structures don't have the same nested structure."):
|
|
structure.to_tensor_list(s_1, value_2)
|
|
|
|
# NOTE(mrry): The repr of the dictionaries is not sorted, so the regexp
|
|
# needs to account for "a" coming before or after "b". It might be worth
|
|
# adding a deterministic repr for these error messages (among other
|
|
# improvements).
|
|
with self.assertRaisesRegex(
|
|
ValueError, "The two structures don't have the same nested structure."):
|
|
structure.to_tensor_list(s_2, value_0)
|
|
|
|
with self.assertRaisesRegex(
|
|
ValueError, "The two structures don't have the same nested structure."):
|
|
structure.to_tensor_list(s_2, value_1)
|
|
|
|
with self.assertRaisesRegex(ValueError,
|
|
r"Cannot create a Tensor from the tensor list"):
|
|
structure.from_tensor_list(s_0, flat_s_1)
|
|
|
|
with self.assertRaisesRegex(ValueError, "Expected 2 tensors but got 3"):
|
|
structure.from_tensor_list(s_0, flat_s_2)
|
|
|
|
with self.assertRaisesRegex(
|
|
ValueError, "Cannot create a SparseTensor from the tensor list"):
|
|
structure.from_tensor_list(s_1, flat_s_0)
|
|
|
|
with self.assertRaisesRegex(ValueError, "Expected 2 tensors but got 3"):
|
|
structure.from_tensor_list(s_1, flat_s_2)
|
|
|
|
with self.assertRaisesRegex(ValueError, "Expected 3 tensors but got 2"):
|
|
structure.from_tensor_list(s_2, flat_s_0)
|
|
|
|
with self.assertRaisesRegex(ValueError, "Expected 3 tensors but got 2"):
|
|
structure.from_tensor_list(s_2, flat_s_1)
|
|
|
|
@combinations.generate(
|
|
combinations.times(test_base.default_test_combinations(),
|
|
_test_convert_legacy_structure_combinations()))
|
|
def testConvertLegacyStructure(self, output_types, output_shapes,
|
|
output_classes, expected_structure):
|
|
actual_structure = structure.convert_legacy_structure(
|
|
output_types, output_shapes, output_classes)
|
|
self.assertEqual(actual_structure, expected_structure)
|
|
|
|
@combinations.generate(test_base.default_test_combinations())
|
|
def testConvertLegacyStructureFail(self):
|
|
with self.assertRaisesRegex(
|
|
TypeError, "Could not build a structure for output class "
|
|
"_EagerTensorArray. Make sure any component class in "
|
|
"`output_classes` inherits from one of the following classes: "
|
|
"`tf.TypeSpec`, `tf.sparse.SparseTensor`, `tf.Tensor`, "
|
|
"`tf.TensorArray`."):
|
|
structure.convert_legacy_structure(dtypes.int32,
|
|
tensor_shape.TensorShape([2, None]),
|
|
tensor_array_ops._EagerTensorArray)
|
|
|
|
@combinations.generate(test_base.default_test_combinations())
|
|
def testNestedNestedStructure(self):
|
|
s = (tensor.TensorSpec([], dtypes.int64),
|
|
(tensor.TensorSpec([], dtypes.float32),
|
|
tensor.TensorSpec([], dtypes.string)))
|
|
|
|
int64_t = constant_op.constant(37, dtype=dtypes.int64)
|
|
float32_t = constant_op.constant(42.0)
|
|
string_t = constant_op.constant("Foo")
|
|
|
|
nested_tensors = (int64_t, (float32_t, string_t))
|
|
|
|
tensor_list = structure.to_tensor_list(s, nested_tensors)
|
|
for expected, actual in zip([int64_t, float32_t, string_t], tensor_list):
|
|
self.assertIs(expected, actual)
|
|
|
|
(actual_int64_t,
|
|
(actual_float32_t,
|
|
actual_string_t)) = structure.from_tensor_list(s, tensor_list)
|
|
self.assertIs(int64_t, actual_int64_t)
|
|
self.assertIs(float32_t, actual_float32_t)
|
|
self.assertIs(string_t, actual_string_t)
|
|
|
|
(actual_int64_t, (actual_float32_t, actual_string_t)) = (
|
|
structure.from_compatible_tensor_list(s, tensor_list))
|
|
self.assertIs(int64_t, actual_int64_t)
|
|
self.assertIs(float32_t, actual_float32_t)
|
|
self.assertIs(string_t, actual_string_t)
|
|
|
|
@combinations.generate(
|
|
combinations.times(test_base.default_test_combinations(),
|
|
_test_batch_combinations()))
|
|
def testBatch(self, element_structure, batch_size,
|
|
expected_batched_structure):
|
|
batched_structure = nest.map_structure(
|
|
lambda component_spec: component_spec._batch(batch_size),
|
|
element_structure)
|
|
self.assertEqual(batched_structure, expected_batched_structure)
|
|
|
|
@combinations.generate(
|
|
combinations.times(test_base.default_test_combinations(),
|
|
_test_unbatch_combinations()))
|
|
def testUnbatch(self, element_structure, expected_unbatched_structure):
|
|
unbatched_structure = nest.map_structure(
|
|
lambda component_spec: component_spec._unbatch(), element_structure)
|
|
self.assertEqual(unbatched_structure, expected_unbatched_structure)
|
|
|
|
# pylint: disable=g-long-lambda
|
|
@combinations.generate(
|
|
combinations.times(test_base.default_test_combinations(),
|
|
_test_to_batched_tensor_list_combinations()))
|
|
def testToBatchedTensorList(self, value_fn, element_0_fn):
|
|
batched_value = value_fn()
|
|
s = structure.type_spec_from_value(batched_value)
|
|
batched_tensor_list = structure.to_batched_tensor_list(s, batched_value)
|
|
|
|
# The batch dimension is 2 for all of the test cases.
|
|
# NOTE(mrry): `tf.shape()` does not currently work for the DT_VARIANT
|
|
# tensors in which we store sparse tensors.
|
|
for t in batched_tensor_list:
|
|
if t.dtype != dtypes.variant:
|
|
self.assertEqual(2, self.evaluate(array_ops.shape(t)[0]))
|
|
|
|
# Test that the 0th element from the unbatched tensor is equal to the
|
|
# expected value.
|
|
expected_element_0 = self.evaluate(element_0_fn())
|
|
unbatched_s = nest.map_structure(
|
|
lambda component_spec: component_spec._unbatch(), s)
|
|
actual_element_0 = structure.from_tensor_list(
|
|
unbatched_s, [t[0] for t in batched_tensor_list])
|
|
|
|
for expected, actual in zip(
|
|
nest.flatten(expected_element_0), nest.flatten(actual_element_0)):
|
|
self.assertValuesEqual(expected, actual)
|
|
|
|
# pylint: enable=g-long-lambda
|
|
|
|
@combinations.generate(test_base.default_test_combinations())
|
|
def testDatasetSpecConstructor(self):
|
|
rt_spec = ragged_tensor.RaggedTensorSpec([10, None], dtypes.int32)
|
|
st_spec = sparse_tensor.SparseTensorSpec([10, 20], dtypes.float32)
|
|
t_spec = tensor.TensorSpec([10, 8], dtypes.string)
|
|
element_spec = {"rt": rt_spec, "st": st_spec, "t": t_spec}
|
|
ds_struct = dataset_ops.DatasetSpec(element_spec, [5])
|
|
self.assertEqual(ds_struct._element_spec, element_spec)
|
|
# Note: shape was automatically converted from a list to a TensorShape.
|
|
self.assertEqual(ds_struct._dataset_shape, tensor_shape.TensorShape([5]))
|
|
|
|
@combinations.generate(test_base.default_test_combinations())
|
|
def testCustomMapping(self):
|
|
elem = CustomMap(foo=constant_op.constant(37.))
|
|
spec = structure.type_spec_from_value(elem)
|
|
self.assertIsInstance(spec, CustomMap)
|
|
self.assertEqual(spec["foo"], tensor.TensorSpec([], dtypes.float32))
|
|
|
|
@combinations.generate(test_base.default_test_combinations())
|
|
def testObjectProxy(self):
|
|
nt_type = collections.namedtuple("A", ["x", "y"])
|
|
proxied = wrapt.ObjectProxy(nt_type(1, 2))
|
|
proxied_spec = structure.type_spec_from_value(proxied)
|
|
self.assertEqual(
|
|
structure.type_spec_from_value(nt_type(1, 2)), proxied_spec)
|
|
|
|
@combinations.generate(test_base.default_test_combinations())
|
|
def testTypeSpecNotBuild(self):
|
|
with self.assertRaisesRegex(
|
|
TypeError, "Could not build a `TypeSpec` for 100 with type int"):
|
|
structure.type_spec_from_value(100, use_fallback=False)
|
|
|
|
@combinations.generate(test_base.default_test_combinations())
|
|
def testTypeSpecNotCompatible(self):
|
|
test_obj = structure.NoneTensorSpec()
|
|
with self.assertRaisesRegex(
|
|
ValueError, r"No `TypeSpec` is compatible with both NoneTensorSpec\(\) "
|
|
"and 100"):
|
|
test_obj.most_specific_compatible_shape(100)
|
|
self.assertEqual(test_obj,
|
|
test_obj.most_specific_compatible_shape(test_obj))
|
|
|
|
@combinations.generate(test_base.default_test_combinations())
|
|
def testDataclasses(self):
|
|
mt = MaskedTensor(mask=True, value=constant_op.constant([1]))
|
|
|
|
mt_type_spec = structure.type_spec_from_value(mt)
|
|
self.assertEqual(mt_type_spec.mask, mt.mask)
|
|
self.assertEqual(
|
|
mt_type_spec.value, structure.type_spec_from_value(mt.value)
|
|
)
|
|
|
|
mt2 = MaskedTensor(mask=True, value=constant_op.constant([2]))
|
|
mt3 = MaskedTensor(mask=False, value=constant_op.constant([1]))
|
|
mt2_type_spec = structure.type_spec_from_value(mt2)
|
|
mt3_type_spec = structure.type_spec_from_value(mt3)
|
|
self.assertEqual(mt_type_spec, mt2_type_spec)
|
|
self.assertNotEqual(mt_type_spec, mt3_type_spec)
|
|
|
|
|
|
class CustomMap(collections_abc.Mapping):
|
|
"""Custom, immutable map."""
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
self.__dict__.update(dict(*args, **kwargs))
|
|
|
|
def __getitem__(self, x):
|
|
return self.__dict__[x]
|
|
|
|
def __iter__(self):
|
|
return iter(self.__dict__)
|
|
|
|
def __len__(self):
|
|
return len(self.__dict__)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test.main()
|