Fixed a bug where the nrows was not correctly set.

The nrows of the child node and the parent node were confused. Since the outcome is immediately modified, this only matters when we update StructuredTensor to have more rigorous unit tests (see cl/456841301).

In addition, added some more helpful debugging so any errors give you information about which field they came from.

PiperOrigin-RevId: 457122359

Author: mzinkevi

struct2tensor/prensor_to_structured_tensor.py

@@ -85,18 +85,21 @@ def _prensor_to_field_map(
     p_fields: Mapping[path.Step, prensor.Prensor],
     nrows: tf.Tensor) -> Mapping[path.Step, structured_tensor.StructuredTensor]:
   """Convert a map of prensors to map of structured tensors."""
-  return {
-      step: _prensor_to_structured_tensor_helper(child, nrows)
-      for step, child in p_fields.items()
-  }
+  result = {}
+  for step, child in p_fields.items():
+    try:
+      result[step] = _prensor_to_structured_tensor_helper(child, nrows)
+    except ValueError as err:
+      raise ValueError(f"Error in field: {step}") from err
+  return result
 
 
 def _child_node_to_structured_tensor(
     node: prensor.ChildNodeTensor, fields: Mapping[path.Step, prensor.Prensor],
     nrows: tf.Tensor) -> structured_tensor.StructuredTensor:
   """Convert a map of prensors to map of structured tensors."""
   st = structured_tensor.StructuredTensor.from_fields(
-      fields=fields, shape=tf.TensorShape([None]), nrows=nrows)
+      fields=fields, shape=tf.TensorShape([None]), nrows=node.size)
   row_partition = RowPartition.from_value_rowids(
       value_rowids=node.parent_index, nrows=nrows)
   return st.partition_outer_dimension(row_partition)

struct2tensor/prensor_to_structured_tensor_test.py

@@ -13,17 +13,15 @@
 # limitations under the License.
 """Tests for StructuredTensor."""
 
+from google.protobuf import text_format
 from struct2tensor import calculate
 from struct2tensor import prensor
 from struct2tensor import prensor_to_structured_tensor
 from struct2tensor.expression_impl import proto
 from struct2tensor.test import prensor_test_util
 from struct2tensor.test import test_pb2
-
 import tensorflow as tf
 
-from google.protobuf import text_format
-
 
 # @test_util.run_all_in_graph_and_eager_modes
 class PrensorToStructuredTensorTest(tf.test.TestCase):