def structured_tensor_to_prensor(
st: structured_tensor.StructuredTensor,
default_field_name: path.Step = "data") -> prensor.Prensor:
"""Converts a structured tensor to a prensor.
Certain rank information must be known. For more details about the
transformation, see the notes above.
Args:
st: the structured tensor to convert.
default_field_name: the name to use when there is an unnamed dimension.
Returns:
a logically equivalent Prensor.
Raises:
ValueError: if there is an issue with the structured tensor.
"""
row_partitions = st.row_partitions
if len(row_partitions) >= 1:
child_prensor = _structured_tensor_to_child_prensor(st, default_field_name)
return prensor.create_prensor_from_root_and_children(
prensor.RootNodeTensor((st).nrows()),
{default_field_name: child_prensor})
elif st.rank == 1:
return prensor.create_prensor_from_root_and_children(
prensor.RootNodeTensor((st).nrows()),
_structured_tensor_prensor_map(st, default_field_name))
else:
# st is a scalar StructuredTensor.
return structured_tensor_to_prensor(_expand_dims(st, 0), default_field_name)
def _one_child_prensor(row_partition: RowPartition,
child_prensor: prensor.Prensor,
default_field_name: path.Step) -> prensor.Prensor:
"""Creates a prensor with a ChildNodeTensor at the root with one child."""
child_node_tensor = _row_partition_to_child_node_tensor(row_partition)
return prensor.create_prensor_from_root_and_children(
child_node_tensor, {default_field_name: child_prensor})
def _to_leaf_prensor_helper(rt: tf.RaggedTensor,
default_field_name: path.Step) -> prensor.Prensor:
"""Converts a fully partitioned ragged tensor to a leaf prensor.
It is assumed that this is a fully partitioned ragged tensor. Specifically,
the values at the end are a vector, not a 2D tensor.
Args:
rt: a fully partitioned ragged tensor (see
_fully_partitioned_ragged_tensor).
default_field_name: a path.Step for unnamed dimensions.
Returns:
a prensor, with a leaf as the root node.
"""
row_partition = rt._row_partition # pylint: disable=protected-access
if rt.ragged_rank == 1:
values = rt.values
leaf = prensor.LeafNodeTensor(row_partition.value_rowids(), values,
True)
return prensor.create_prensor_from_root_and_children(leaf, {})
else:
return _one_child_prensor(
row_partition,
_to_leaf_prensor_helper(rt.values, default_field_name),
default_field_name)
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)
def test_non_root_prensor(self):
child_prensor = prensor.create_prensor_from_root_and_children(
prensor_test_util.create_child_node([0, 0, 1, 3, 7], True), {})
with self.assertRaisesRegexp(ValueError, "Must be a root prensor"):
prensor_to_structured_tensor.prensor_to_structured_tensor(
child_prensor)