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)