def testNestedStructConstruction(self):
rt = ragged_factory_ops.constant([[1, 2], [3]])
struct1 = StructuredTensor.from_fields(shape=[], fields={"x": [1, 2]})
struct2 = StructuredTensor.from_fields(shape=[2], fields={"x": [1, 2]})
struct3 = StructuredTensor.from_fields(shape=[],
fields={
"r": rt,
"s": struct1
})
struct4 = StructuredTensor.from_fields(shape=[2],
fields={
"r": rt,
"s": struct2
})
self.assertEqual(struct3.shape.as_list(), [])
self.assertEqual(struct3.rank, 0)
self.assertEqual(set(struct3.field_names()), set(["r", "s"]))
self.assertAllEqual(struct3.field_value("r"), rt)
self.assertAllEqual(struct3.field_value("s"), struct1)
self.assertEqual(struct4.shape.as_list(), [2])
self.assertEqual(struct4.rank, 1)
self.assertEqual(set(struct4.field_names()), set(["r", "s"]))
self.assertAllEqual(struct4.field_value("r"), rt)
self.assertAllEqual(struct4.field_value("s"), struct2)
def _structured_tensor_prensor_map(
st: structured_tensor.StructuredTensor,
default_field_name: path.Step) -> Mapping[path.Step, prensor.Prensor]:
"""Creates a map of fields, to put in a child or root prensor."""
return {
k: _structured_tensor_field_to_prensor(
st.field_value(k), default_field_name) for k in st.field_names()
}
def testToFromComponents(self, shape, fields, field_specs):
struct = StructuredTensor.from_fields(fields, shape)
spec = StructuredTensor.Spec(_ragged_shape=DynamicRaggedShape.Spec(
row_partitions=[], static_inner_shape=shape, dtype=dtypes.int64),
_fields=field_specs)
actual_components = spec._to_components(struct)
rt_reconstructed = spec._from_components(actual_components)
self.assertAllEqual(struct, rt_reconstructed)
def testConcatTuple(self):
values = (StructuredTensor.from_pyval([{
"a": 3
}]), StructuredTensor.from_pyval([{
"a": 4
}]))
actual = array_ops.concat(values, axis=0)
self.assertAllEqual(actual, [{"a": 3}, {"a": 4}])
def _expand_dims_scalar(st: structured_tensor.StructuredTensor):
"""_expand_dims for a scalar structured tensor."""
new_shape = tf.constant([1], dtype=tf.int64)
new_fields = {
k: _expand_dims(st.field_value(k), 0)
for k in st.field_names()
}
return structured_tensor.StructuredTensor.from_fields(new_fields,
shape=new_shape)
def testRandomShuffle2022Eager(self):
original = StructuredTensor.from_pyval([{
"x0": 0,
"y": {
"z": [[3, 13]]
}
}, {
"x0": 1,
"y": {
"z": [[3], [4, 13]]
}
}, {
"x0": 2,
"y": {
"z": [[3, 5], [4]]
}
}, {
"x0": 3,
"y": {
"z": [[3, 7, 1], [4]]
}
}, {
"x0": 4,
"y": {
"z": [[3], [4]]
}
}]) # pyformat: disable
expected = StructuredTensor.from_pyval([{
"x0": 1,
"y": {
"z": [[3], [4, 13]]
}
}, {
"x0": 0,
"y": {
"z": [[3, 13]]
}
}, {
"x0": 3,
"y": {
"z": [[3, 7, 1], [4]]
}
}, {
"x0": 4,
"y": {
"z": [[3], [4]]
}
}, {
"x0": 2,
"y": {
"z": [[3, 5], [4]]
}
}]) # pyformat: disable
random_seed.set_seed(1066)
result = structured_array_ops.random_shuffle(original, seed=2022)
self.assertAllEqual(result, expected)
def _structured_tensor_like(t):
"""Create a StructuredTensor with the shape of a (composite) tensor."""
if isinstance(t, ops.Tensor):
return _structured_tensor_from_dense_tensor(t)
if ragged_tensor.is_ragged(t):
return StructuredTensor.from_fields(
{}, shape=t.get_shape(), row_partitions=_all_nested_row_partitions(t))
# here, it is a StructuredTensor
return StructuredTensor.from_fields({},
shape=t.shape,
row_partitions=t.row_partitions,
nrows=t.nrows())
def _structured_tensor_from_dense_tensor(t):
"""Create a structured tensor with the shape of a dense tensor."""
# Note: If a tensor will have rank 0,
# it either has a fully defined shape or has unknown rank.
if t.shape.is_fully_defined():
return StructuredTensor.from_fields({}, shape=t.shape)
elif t.shape.rank is None:
raise ValueError("Can't build StructuredTensor w/ unknown rank")
elif t.shape.rank == 1:
return StructuredTensor.from_fields({}, shape=t.shape,
nrows=array_ops.shape(t)[0])
else:
rt = ragged_tensor.RaggedTensor.from_tensor(t)
return _structured_tensor_from_row_partitions(t.shape,
rt._nested_row_partitions)
def testExtendOpErrorNotList(self):
# Should be a list.
values = StructuredTensor.from_pyval({})
def leaf_op(values):
return values[0]
with self.assertRaisesRegex(ValueError, "Expected a list"):
structured_array_ops._extend_op(values, leaf_op)
def testToFromComponentsEmptyScalar(self):
struct = StructuredTensor.from_fields(fields={}, shape=[])
spec = struct._type_spec
components = spec._to_components(struct)
rt_reconstructed = spec._from_components(components)
self.assertAllEqual(struct, rt_reconstructed)
self.assertEqual(components, ({}, (), ()))
def testFromFields(self,
shape,
fields,
expected_shape=None,
nrows=None,
row_partitions=None):
if callable(fields):
fields = fields(
) # deferred construction: fields may include tensors.
if callable(nrows):
nrows = nrows() # deferred construction.
if callable(row_partitions):
row_partitions = row_partitions() # deferred construction.
for validate in (True, False):
struct = StructuredTensor.from_fields(
fields,
shape,
nrows=nrows,
row_partitions=row_partitions,
validate=validate)
if expected_shape is None:
expected_shape = shape
self.assertEqual(struct.shape.as_list(), expected_shape)
self.assertLen(expected_shape, struct.rank)
self.assertCountEqual(struct.field_names(), tuple(fields.keys()))
for field, value in fields.items():
self.assertIsInstance(
struct.field_value(field),
(ops.Tensor, structured_tensor.StructuredTensor,
ragged_tensor.RaggedTensor))
self.assertAllEqual(struct.field_value(field), value)
def testPartitionOuterDims(self):
if not context.executing_eagerly(): return # TESTING
a = dict(x=1, y=[1, 2])
b = dict(x=2, y=[3, 4])
c = dict(x=3, y=[5, 6])
d = dict(x=4, y=[7, 8])
st1 = StructuredTensor.from_pyval([a, b, c, d])
st2 = st1.partition_outer_dimension(
row_partition.RowPartition.from_row_splits([0, 2, 2, 3, 4]))
self.assertAllEqual(st2, [[a, b], [], [c], [d]])
st3 = st2.partition_outer_dimension(
row_partition.RowPartition.from_row_lengths([1, 0, 3, 0]))
self.assertAllEqual(st3, [[[a, b]], [], [[], [c], [d]], []])
# If we partition with uniform_row_lengths, then `x` is partitioned into
# a Tensor (not a RaggedTensor).
st4 = st1.partition_outer_dimension(
row_partition.RowPartition.from_uniform_row_length(
uniform_row_length=2, nvals=4, nrows=2))
self.assertAllEqual(
st4,
structured_tensor.StructuredTensor.from_pyval(
[[a, b], [c, d]],
structured_tensor.StructuredTensorSpec(
[2, 2], {
"x":
tensor_spec.TensorSpec([2, 2], dtypes.int32),
"y":
ragged_tensor.RaggedTensorSpec([2, 2, None],
dtypes.int32)
})))
def testTupleFieldValue(self):
st = StructuredTensor.from_pyval({"a": 5, "b": {"c": [1, 2, 3]}})
self.assertAllEqual(st.field_value(("a", )), 5)
self.assertAllEqual(st.field_value(("b", "c")), [1, 2, 3])
expected = "Field path \(.*a.*,.*b.*\) not found in .*"
with self.assertRaisesRegexp(KeyError, expected):
st.field_value(("a", "b"))
def testTupleFieldValue(self):
st = StructuredTensor.from_pyval({"a": 5, "b": {"c": [1, 2, 3]}})
self.assertAllEqual(st.field_value(("a", )), 5)
self.assertAllEqual(st.field_value(("b", "c")), [1, 2, 3])
with self.assertRaisesRegexp(
KeyError, r"Field path \('a', 'b'\) not found in .*"):
st.field_value(("a", "b"))
def testFromFieldsErrors(self,
fields,
shape,
nrows=None,
row_partitions=None,
validate=False,
err=ValueError,
msg=None,
test_in_eager=True):
if not test_in_eager and context.executing_eagerly():
return
if callable(fields):
fields = fields() # deferred construction.
if callable(nrows):
nrows = nrows() # deferred construction.
if callable(row_partitions):
row_partitions = row_partitions() # deferred construction.
with self.assertRaisesRegexp(err, msg):
struct = StructuredTensor.from_fields(
fields=fields,
shape=shape,
nrows=nrows,
row_partitions=row_partitions,
validate=validate)
for field_name in struct.field_names():
self.evaluate(struct.field_value(field_name))
self.evaluate(struct.nrows())
def testToFromComponents(self, shape, fields, field_specs):
struct = StructuredTensor.from_fields(fields, shape)
spec = StructuredTensorSpec(shape, field_specs)
actual_components = spec._to_components(struct)
self.assertLen(actual_components, 3)
self.assertAllTensorsEqual(actual_components[0], fields)
rt_reconstructed = spec._from_components(actual_components)
self.assertAllEqual(struct, rt_reconstructed)
def testExpandDimsScalar(self):
# Note that if we expand_dims for the final dimension and there are scalar
# fields, then the shape is (2, None, None, 1), whereas if it is constructed
# from pyval it is (2, None, None, None).
st = [[[{"x": 1}, {"x": 2}], [{"x": 3}]], [[{"x": 4}]]]
st = StructuredTensor.from_pyval(st)
result = array_ops.expand_dims(st, 3)
expected_shape = tensor_shape.TensorShape([2, None, None, 1])
self.assertEqual(repr(expected_shape), repr(result.shape))
def testSizeAlt(self, values, dtype, expected):
st = StructuredTensor.from_pyval(values)
# NOTE: size is very robust. There aren't arguments that
# should cause this operation to fail.
actual = array_ops.size(st, out_type=dtype)
self.assertAllEqual(actual, expected)
actual2 = array_ops.size_v2(st, out_type=dtype)
self.assertAllEqual(actual2, expected)
def testOnesLikeObjectAlt(self, values, dtype, expected):
st = StructuredTensor.from_pyval(values)
# NOTE: ones_like is very robust. There aren't arguments that
# should cause this operation to fail.
actual = array_ops.ones_like(st, dtype)
self.assertAllEqual(actual, expected)
actual2 = array_ops.ones_like_v2(st, dtype)
self.assertAllEqual(actual2, expected)
def testConstruction(self):
spec1_fields = dict(a=T_1_2_3_4)
spec1 = StructuredTensor.Spec(_ragged_shape=DynamicRaggedShape.Spec(
row_partitions=[],
static_inner_shape=tensor_shape.TensorShape([1, 2, 3]),
dtype=dtypes.int64),
_fields=spec1_fields)
self.assertEqual(spec1._shape, (1, 2, 3))
self.assertEqual(spec1._field_specs, spec1_fields)
spec2_fields = dict(a=T_1_2, b=T_1_2_8, c=R_1_N, d=R_1_N_N, s=spec1)
spec2 = StructuredTensor.Spec(_ragged_shape=DynamicRaggedShape.Spec(
row_partitions=[],
static_inner_shape=tensor_shape.TensorShape([1, 2]),
dtype=dtypes.int64),
_fields=spec2_fields)
self.assertEqual(spec2._shape, (1, 2))
self.assertEqual(spec2._field_specs, spec2_fields)
def testPartitionOuterDimsErrors(self):
st = StructuredTensor.from_fields({})
partition = row_partition.RowPartition.from_row_splits([0])
with self.assertRaisesRegexp(ValueError,
r"Shape \(\) must have rank at least 1"):
st.partition_outer_dimension(partition)
with self.assertRaisesRegexp(TypeError,
"row_partition must be a RowPartition"):
st.partition_outer_dimension(10)
def testGather(self, params, indices, axis, batch_dims, expected):
params = StructuredTensor.from_pyval(params)
# validate_indices isn't actually used, and we aren't testing names
actual = array_ops.gather(params,
indices,
validate_indices=True,
axis=axis,
name=None,
batch_dims=batch_dims)
self.assertAllEqual(actual, expected)
def testRandomShuffleScalarError(self):
original = StructuredTensor.from_pyval({
"x0": 2,
"y": {
"z": [[3, 5], [4]]
}
}) # pyformat: disable
with self.assertRaisesRegex(ValueError, "scalar"):
random_ops.random_shuffle(original)
def testGatherError(self, params, indices, axis, batch_dims, error_type,
error_regex):
params = StructuredTensor.from_pyval(params)
with self.assertRaisesRegex(error_type, error_regex):
structured_array_ops.gather(params,
indices,
validate_indices=True,
axis=axis,
name=None,
batch_dims=batch_dims)
def testMergeDims_0_1(self):
rt = ragged_tensor.RaggedTensor.from_value_rowids(
array_ops.constant([[1, 2], [3, 4], [5, 6]]), [0, 0, 1])
struct = StructuredTensor.from_fields({"r": rt}, [2])
struct_2 = struct.partition_outer_dimension(
row_partition.RowPartition.from_row_splits([0, 1, 2]))
struct_3 = struct_2.partition_outer_dimension(
row_partition.RowPartition.from_row_splits([0, 1, 2]))
self.assertLen(struct_3.row_partitions, 2)
merged = struct_3.merge_dims(0, 1)
self.assertLen(merged.row_partitions, 1)
def testGatherRagged(self, params, indices, axis, batch_dims, expected):
params = StructuredTensor.from_pyval(params)
# Shouldn't need to do this, but see cl/366396997
indices = ragged_factory_ops.constant(indices)
# validate_indices isn't actually used, and we aren't testing names
actual = array_ops.gather(params,
indices,
validate_indices=True,
axis=axis,
name=None,
batch_dims=batch_dims)
self.assertAllEqual(actual, expected)
def testRank(self, row_partitions, shape, expected):
if row_partitions is not None:
row_partitions = [
row_partition.RowPartition.from_row_splits(r) for r in row_partitions
]
st = StructuredTensor.from_fields({},
shape=shape,
row_partitions=row_partitions)
# NOTE: rank is very robust. There aren't arguments that
# should cause this operation to fail.
actual = structured_array_ops.rank(st)
self.assertAllEqual(expected, actual)
def testZerosLikeObject(self, row_partitions, shape, dtype, expected):
if row_partitions is not None:
row_partitions = [
row_partition.RowPartition.from_row_splits(r)
for r in row_partitions
]
st = StructuredTensor.from_fields({},
shape=shape,
row_partitions=row_partitions)
# NOTE: zeros_like is very robust. There aren't arguments that
# should cause this operation to fail.
actual = structured_array_ops.zeros_like_object(st, dtype)
self.assertAllEqual(actual, expected)
def _lambda_for_fields(self):
return lambda: {
'a':
np.ones([1, 2, 3, 1]),
'b':
np.ones([1, 2, 3, 1, 5]),
'c':
ragged_factory_ops.constant(np.zeros([1, 2, 3, 1], dtype=np.uint8),
dtype=dtypes.uint8),
'd':
ragged_factory_ops.constant(np.zeros([1, 2, 3, 1, 3]).tolist(),
ragged_rank=1),
'e':
ragged_factory_ops.constant(np.zeros([1, 2, 3, 1, 2, 2]).tolist(),
ragged_rank=2),
'f':
ragged_factory_ops.constant(np.zeros([1, 2, 3, 1, 3]),
dtype=dtypes.float32),
'g':
StructuredTensor.from_pyval([[
[ # pylint: disable=g-complex-comprehension
[{
'x': j,
'y': k
}] for k in range(3)
] for j in range(2)
]]),
'h':
StructuredTensor.from_pyval([[
[ # pylint: disable=g-complex-comprehension
[[{
'x': j,
'y': k,
'z': z
} for z in range(j)]] for k in range(3)
] for j in range(2)
]]),
}
def _structured_tensor_to_child_prensor(
st: structured_tensor.StructuredTensor,
default_field_name: path.Step) -> prensor.Prensor:
"""Creates a prensor with a ChildNodeTensor at the root."""
row_partitions = st.row_partitions
if len(row_partitions) == 1:
child_st = st.merge_dims(0, 1)
row_partition = row_partitions[0]
return prensor.create_prensor_from_root_and_children(
_row_partition_to_child_node_tensor(row_partition),
_structured_tensor_prensor_map(child_st, default_field_name))
elif len(row_partitions) > 1:
row_partition = row_partitions[0]
child_st = st.merge_dims(0, 1)
return _one_child_prensor(
row_partition,
_structured_tensor_to_child_prensor(child_st, default_field_name),
default_field_name)
# This requires us to transform the scalar to a vector.
# The fields could be scalars or vectors.
# We need _expand_dims(...) to make this work.
return _structured_tensor_to_child_prensor(
_expand_dims(st, 1), default_field_name)
