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)