def to_sparse_tensor_value( sparse_tensor: Union[tf.SparseTensor, tf.compat.v1.SparseTensorValue] ) -> types.SparseTensorValue: """Converts sparse tensor to SparseTensorValue.""" return types.SparseTensorValue(to_numpy(sparse_tensor.values), to_numpy(sparse_tensor.indices), to_numpy(sparse_tensor.dense_shape))
def testToFromTensorValues(self): tensor_values = { 'features': { 'feature_1': np.array([1, 2, 3]), 'feature_2': types.SparseTensorValue(values=np.array([0.5, -1., 0.5, -1.]), indices=np.array([[0, 3, 1], [0, 20, 0], [1, 3, 1], [1, 20, 0]]), dense_shape=np.array([2, 100, 3])), 'feature_3': types.RaggedTensorValue( values=np.array([3, 1, 4, 1, 5, 9, 2, 7, 1, 8, 8, 2, 1]), nested_row_splits=[ np.array([0, 3, 6]), np.array([0, 2, 3, 4, 5, 5, 8]), np.array([0, 2, 3, 3, 6, 9, 10, 11, 13]) ]), 'feature_4': types.VarLenTensorValue(values=np.array([1, 2, 3]), indices=np.array([[0, 0], [0, 1], [1, 0]]), dense_shape=np.array([2, 2])) }, 'labels': np.array([1]) } actual = util.to_tensor_values( util.to_tensorflow_tensors(tensor_values)) self.assertAllClose(actual, tensor_values)
def testToScalar(self): self.assertEqual(1, metric_util.to_scalar(np.array([1]))) self.assertEqual(1.0, metric_util.to_scalar(np.array(1.0))) self.assertEqual('string', metric_util.to_scalar(np.array([['string']]))) sparse_tensor = types.SparseTensorValue(indices=np.array([0]), values=np.array([1]), dense_shape=np.array([1])) self.assertEqual(1, metric_util.to_scalar(sparse_tensor))
def testPrepareLabelsAndPredictionsSparseTensorValueAndVocab(self): labels = types.SparseTensorValue(indices=np.array([0, 2]), values=np.array(['c', 'a']), dense_shape=np.array([1, 2])) preds = { 'probabilities': [0.2, 0.7, 0.1], 'all_classes': ['a', 'b', 'c'] } got_labels, got_preds = metric_util.prepare_labels_and_predictions( labels, preds) self.assertAllClose(got_labels, np.array([1, 0, 1])) self.assertAllClose(got_preds, np.array([0.2, 0.7, 0.1]))
def testToTFSparseTensorFromSparseTensorValue(self): original = types.SparseTensorValue(values=np.array( [0.5, -1., 0.5, -1.]), indices=np.array([[0, 3, 1], [0, 20, 0], [1, 3, 1], [1, 20, 0]]), dense_shape=np.array([2, 100, 3])) sparse_tensor = util.to_tensorflow_tensor(original) self.assertAllClose(sparse_tensor.values.numpy(), original.values) self.assertAllClose(sparse_tensor.indices.numpy(), original.indices) self.assertAllClose(sparse_tensor.dense_shape.numpy(), original.dense_shape)
def testStandardMetricInputsWithSparseTensorValue(self): example = metric_types.StandardMetricInputs( labels=types.SparseTensorValue(values=np.array([1]), indices=np.array([2]), dense_shape=np.array([0, 1])), predictions=np.array([0, 0.5, 0.3, 0.9]), example_weights=np.array([0.0])) iterator = metric_util.to_label_prediction_example_weight(example) for expected_label, expected_prediction in zip((0.0, 0.0, 1.0, 0.0), (0.0, 0.5, 0.3, 0.9)): got_label, got_pred, got_example_weight = next(iterator) self.assertAllClose(got_label, np.array([expected_label])) self.assertAllClose(got_pred, np.array([expected_prediction])) self.assertAllClose(got_example_weight, np.array([0.0]))
def testMergeExtractsRaisesException(self): extracts = [ { 'features': { 'feature_3': types.SparseTensorValue(values=np.array([1]), indices=np.array([[0, 1]]), dense_shape=np.array([1, 2])) }, }, { 'features': { 'feature_3': types.SparseTensorValue(values=np.array([2]), indices=np.array([[0, 2]]), dense_shape=np.array([1, 3])) }, }, ] with self.assertRaisesWithPredicateMatch( RuntimeError, lambda exc: isinstance(exc.__cause__, RuntimeError)): util.merge_extracts(extracts)
def testToFromTensorValuesWithSpecs(self): sparse_value = types.SparseTensorValue( values=np.array([0.5, -1., 0.5, -1.], dtype=np.float32), indices=np.array([[0, 3, 1], [0, 20, 0], [1, 3, 1], [1, 20, 0]]), dense_shape=np.array([2, 100, 3])) ragged_value = types.RaggedTensorValue( values=np.array([3, 1, 4, 1, 5, 9, 2, 7, 1, 8, 8, 2, 1], dtype=np.float32), nested_row_splits=[ np.array([0, 3, 6]), np.array([0, 2, 3, 4, 5, 5, 8]), np.array([0, 2, 3, 3, 6, 9, 10, 11, 13]) ]) tensor_values = { 'features': { 'feature_1': np.array([1, 2, 3], dtype=np.float32), 'feature_2': sparse_value, 'feature_3': ragged_value, 'ignored_feature': np.array([1, 2, 3]) }, 'labels': np.array([1], dtype=np.float32), 'ignored': np.array([2]) } specs = { 'features': { 'feature_1': tf.TensorSpec([3], dtype=tf.float32), 'feature_2': tf.SparseTensorSpec(shape=[2, 100, 3], dtype=tf.float32), 'feature_3': tf.RaggedTensorSpec(shape=[2, None, None, None], dtype=tf.float32) }, 'labels': tf.TensorSpec([1], dtype=tf.float32) } actual = util.to_tensor_values( util.to_tensorflow_tensors(tensor_values, specs)) expected = { 'features': { 'feature_1': np.array([1, 2, 3], dtype=np.float32), 'feature_2': sparse_value, 'feature_3': ragged_value }, 'labels': np.array([1], dtype=np.float32) } self.assertAllClose(actual, expected)
def testBatchSizeWithTensorValues(self): tensor_values = { 'feature_1': np.array([1, 2], dtype=np.float32), 'feature_2': types.SparseTensorValue(values=np.array([0.5, -1., 0.5, -1.], dtype=np.float32), indices=np.array([[0, 3, 1], [0, 20, 0], [1, 3, 1], [1, 20, 0]]), dense_shape=np.array([2, 100, 3])), 'feature_3': types.RaggedTensorValue( values=np.array([3, 1, 4, 1, 5, 9, 2, 7, 1, 8, 8, 2, 1], dtype=np.float32), nested_row_splits=[ np.array([0, 3, 6]), np.array([0, 2, 3, 4, 5, 5, 8]), np.array([0, 2, 3, 3, 6, 9, 10, 11, 13]) ]), } self.assertEqual(util.batch_size(tensor_values), 2)
def _split_sparse_batch( sparse_batch: types.SparseTensorValue ) -> List[types.SparseTensorValue]: """Converts a batched SparseTensorValue to a list of SparseTensorValues.""" result = [] batch_indices = sparse_batch.indices[:, 0] split_shape = sparse_batch.dense_shape[1:] for i in range(sparse_batch.dense_shape[0]): element_mask = batch_indices == i split_batch_indices = sparse_batch.indices[element_mask, :] split_indices = split_batch_indices[:, 1:] # Note that this style of indexing will create a copy rather than a view of # the underlying array. Consider updating SparseTensorValues to be more # strict about the layout of values, and use simple indexing (a[start:end]) # to extract contiguous ranges of values and indices, similar to how # VarLenTensorValue.dense_rows() is implemented. split_values = sparse_batch.values[element_mask] result.append( types.SparseTensorValue(values=split_values, indices=split_indices, dense_shape=split_shape)) return result
def testInferTensorSpecs(self): sparse_value = types.SparseTensorValue( values=np.array([0.5, -1., 0.5, -1.], dtype=np.float32), indices=np.array([[0, 3, 1], [0, 20, 0], [1, 3, 1], [1, 20, 0]]), dense_shape=np.array([2, 100, 3])) ragged_value = types.RaggedTensorValue( values=np.array([3, 1, 4, 1, 5, 9, 2, 7, 1, 8, 8, 2, 1], dtype=np.float32), nested_row_splits=[ np.array([0, 3, 6]), np.array([0, 2, 3, 4, 5, 5, 8]), np.array([0, 2, 3, 3, 6, 9, 10, 11, 13]) ]) tensor_values = { 'features': { 'feature_1': np.array([1, 2, 3], dtype=np.float32), 'feature_2': sparse_value, 'feature_3': ragged_value, }, 'labels': np.array([1], dtype=np.float32), } expected_specs = { 'features': { 'feature_1': tf.TensorSpec([None], dtype=tf.float32), 'feature_2': tf.SparseTensorSpec(shape=[None, 100, 3], dtype=tf.float32), 'feature_3': tf.RaggedTensorSpec(shape=[None, None, None, None], dtype=tf.float32) }, 'labels': tf.TensorSpec([None], dtype=tf.float32) } got_specs = util.infer_tensor_specs( util.to_tensorflow_tensors(tensor_values)) self.assertDictEqual(expected_specs, got_specs)
def testSplitExtracts(self): extracts = { 'features': { 'feature_1': np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]), 'feature_2': np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]), 'feature_3': types.SparseTensorValue(values=np.array([1, 2, 3]), indices=np.array([[0, 0, 1], [1, 0, 2], [2, 0, 0]]), dense_shape=np.array([3, 1, 3])), 'feature_4': types.RaggedTensorValue(values=np.array([ 3, 1, 4, 1, 5, 9, 2, 6, 3, 1, 4, 1, 5, 9, 2, 6, 3, 1, 4, 1, 5, 9, 2, 6 ]), nested_row_splits=[ np.array([0, 5, 10, 15]), np.array([ 0, 4, 4, 7, 8, 8, 12, 12, 15, 16, 16, 20, 20, 23, 24, 24 ]) ]), 'feature_5': types.VarLenTensorValue(values=np.array([1, 2, 3, 3, 3]), indices=np.array([[0, 0], [1, 0], [2, 0], [2, 1], [2, 2]]), dense_shape=np.array([3, 3])), 'feature_6': types.VarLenTensorValue(values=np.array([1, 3, 3, 3]), indices=np.array([[0, 0], [2, 0], [2, 1], [2, 2]]), dense_shape=np.array([3, 3])), }, 'labels': np.array([1.0, 0.0, 1.0]), 'example_weights': np.array([0.0, 0.5, 1.0]), 'predictions': { 'model1': np.array([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]]), 'model2': np.array([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]]) }, 'empty': None, 'multi_level_empty': { 'empty': None, 'next_level': { 'empty': None }, }, '_slice_key_types': types.VarLenTensorValue.from_dense_rows([ slicer_lib.slice_keys_to_numpy_array([(('gender', 'm'), ), ()]), slicer_lib.slice_keys_to_numpy_array([()]), slicer_lib.slice_keys_to_numpy_array([()]) ]) } expected = [ { 'features': { 'feature_1': np.array([1.0, 2.0]), 'feature_2': np.array([1.0, 2.0]), 'feature_3': types.SparseTensorValue(values=np.array([1]), indices=np.array([[0, 1]]), dense_shape=np.array([1, 3])), 'feature_4': types.RaggedTensorValue( values=np.array([3, 1, 4, 1, 5, 9, 2, 6]), nested_row_splits=[np.array([0, 4, 4, 7, 8, 8])]), 'feature_5': np.array([1.0]), 'feature_6': np.array([1.0]), }, 'labels': np.array([1.0]), 'example_weights': np.array([0.0]), 'predictions': { 'model1': np.array([0.1, 0.2]), 'model2': np.array([0.1, 0.2]) }, 'empty': None, 'multi_level_empty': { 'empty': None, 'next_level': { 'empty': None }, }, '_slice_key_types': np.array([(('gender', 'm'), ), ()], dtype=object) }, { 'features': { 'feature_1': np.array([3.0, 4.0]), 'feature_2': np.array([3.0, 4.0]), 'feature_3': types.SparseTensorValue(values=np.array([2]), indices=np.array([[0, 2]]), dense_shape=np.array([1, 3])), 'feature_4': types.RaggedTensorValue( values=np.array([3, 1, 4, 1, 5, 9, 2, 6]), nested_row_splits=[np.array([0, 4, 4, 7, 8, 8])]), 'feature_5': np.array([2.0]), 'feature_6': np.array([]), }, 'labels': np.array([0.0]), 'example_weights': np.array([0.5]), 'predictions': { 'model1': np.array([0.3, 0.4]), 'model2': np.array([0.3, 0.4]) }, 'empty': None, 'multi_level_empty': { 'empty': None, 'next_level': { 'empty': None }, }, '_slice_key_types': slicer_lib.slice_keys_to_numpy_array([()]) }, { 'features': { 'feature_1': np.array([5.0, 6.0]), 'feature_2': np.array([5.0, 6.0]), 'feature_3': types.SparseTensorValue(values=np.array([3]), indices=np.array([[0, 0]]), dense_shape=np.array([1, 3])), 'feature_4': types.RaggedTensorValue( values=np.array([3, 1, 4, 1, 5, 9, 2, 6]), nested_row_splits=[np.array([0, 4, 4, 7, 8, 8])]), 'feature_5': np.array([3.0, 3.0, 3.0]), 'feature_6': np.array([3.0, 3.0, 3.0]) }, 'labels': np.array([1.0]), 'example_weights': np.array([1.0]), 'predictions': { 'model1': np.array([0.5, 0.6]), 'model2': np.array([0.5, 0.6]) }, 'empty': None, 'multi_level_empty': { 'empty': None, 'next_level': { 'empty': None }, }, '_slice_key_types': slicer_lib.slice_keys_to_numpy_array([()]) }, ] np.testing.assert_equal(util.split_extracts(extracts), expected)
def testMergeExtracts(self): extracts = [ { 'features': { 'feature_1': np.array([1.0, 2.0]), 'feature_2': np.array([1.0, 2.0]), 'feature_3': types.SparseTensorValue(values=np.array([1]), indices=np.array([[0, 1]]), dense_shape=np.array([1, 3])), 'feature_4': types.RaggedTensorValue( values=np.array([3, 1, 4, 1, 5, 9, 2, 6]), nested_row_splits=[np.array([0, 4, 4, 7, 8, 8])]), 'feature_5': types.SparseTensorValue(values=np.array([1]), indices=np.array([[0, 1]]), dense_shape=np.array([1, 3])), 'feature_6': np.array([]), }, 'labels': np.array([1.0]), 'example_weights': np.array(0.0), 'predictions': { 'model1': np.array([0.1, 0.2]), 'model2': np.array([0.1, 0.2]) }, '_slice_key_types': slicer_lib.slice_keys_to_numpy_array([('gender', 'm'), ()]) }, { 'features': { 'feature_1': np.array([3.0, 4.0]), 'feature_2': np.array([3.0, 4.0]), 'feature_3': types.SparseTensorValue(values=np.array([2]), indices=np.array([[0, 2]]), dense_shape=np.array([1, 3])), 'feature_4': types.RaggedTensorValue( values=np.array([3, 1, 4, 1, 5, 9, 2, 6]), nested_row_splits=[np.array([0, 4, 4, 7, 8, 8])]), 'feature_5': types.SparseTensorValue(values=np.array([2]), indices=np.array([[0, 2]]), dense_shape=np.array([1, 4])), 'feature_6': np.array([]), }, 'labels': np.array([0.0]), 'example_weights': np.array(0.5), 'predictions': { 'model1': np.array([0.3, 0.4]), 'model2': np.array([0.3, 0.4]) }, '_slice_key_types': slicer_lib.slice_keys_to_numpy_array([()]) }, { 'features': { 'feature_1': np.array([5.0, 6.0]), 'feature_2': np.array([5.0, 6.0]), 'feature_3': types.SparseTensorValue(values=np.array([3]), indices=np.array([[0, 0]]), dense_shape=np.array([1, 3])), 'feature_4': types.RaggedTensorValue( values=np.array([3, 1, 4, 1, 5, 9, 2, 6]), nested_row_splits=[np.array([0, 4, 4, 7, 8, 8])]), 'feature_5': types.SparseTensorValue(values=np.array([3]), indices=np.array([[0, 3]]), dense_shape=np.array([1, 5])), 'feature_6': np.array([2.0, 3.0]), }, 'labels': np.array([1.0]), 'example_weights': np.array(1.0), 'predictions': { 'model1': np.array([0.5, 0.6]), 'model2': np.array([0.5, 0.6]) }, '_slice_key_types': slicer_lib.slice_keys_to_numpy_array([()]) }, ] expected = { 'features': { 'feature_1': np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]), 'feature_2': np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]), 'feature_3': types.SparseTensorValue(values=np.array([1, 2, 3]), indices=np.array([[0, 0, 1], [1, 0, 2], [2, 0, 0]]), dense_shape=np.array([3, 1, 3])), 'feature_4': types.RaggedTensorValue(values=np.array([ 3, 1, 4, 1, 5, 9, 2, 6, 3, 1, 4, 1, 5, 9, 2, 6, 3, 1, 4, 1, 5, 9, 2, 6 ]), nested_row_splits=[ np.array([0, 5, 10, 15]), np.array([ 0, 4, 4, 7, 8, 8, 12, 12, 15, 16, 16, 20, 20, 23, 24, 24 ]) ]), 'feature_5': types.SparseTensorValue(values=np.array([1, 2, 3]), indices=np.array([[0, 0, 1], [1, 0, 2], [2, 0, 3]]), dense_shape=np.array([3, 1, 5])), 'feature_6': types.VarLenTensorValue(values=np.array([2.0, 3.0]), indices=np.array([[2, 0], [2, 1]]), dense_shape=np.array([3, 2])) }, 'labels': np.array([1.0, 0.0, 1.0]), 'example_weights': np.array([0.0, 0.5, 1.0]), 'predictions': { 'model1': np.array([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]]), 'model2': np.array([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]]) }, '_slice_key_types': types.VarLenTensorValue( values=slicer_lib.slice_keys_to_numpy_array([('gender', 'm'), (), (), ()]), indices=np.array([[0, 0], [0, 1], [1, 0], [2, 0]]), dense_shape=np.array([3, 2])) } np.testing.assert_equal(util.merge_extracts(extracts), expected)
class SliceAccessorTest(tf.test.TestCase, parameterized.TestCase): def testRaisesKeyError(self): accessor = slice_accessor.SliceAccessor({}) with self.assertRaises(KeyError): accessor.get('no_such_key') @parameterized.named_parameters( ('sparse_tensor_value', types.SparseTensorValue(indices=np.array([[0, 0], [1, 1]]), values=np.array(['apple', 'banana']), dense_shape=np.array( [2, 2])), ['apple', 'banana']), ('ragged_tensor_value', types.RaggedTensorValue(values=np.array([1, 2, 3]), nested_row_splits=[np.array([0, 0, 1])]), [1, 2, 3]), ('dense', np.array([1.0, 2.0]), [1.0, 2.0]), ('dense_single', np.array([7.0]), [7.0]), ('dense_multidim', np.array([[1.0, 2.0], [3.0, 4.0]]), [1.0, 2.0, 3.0, 4.0]), ('squeeze_needed', np.array([[2.0]]), [2.0]), ('list', [1, 2, 3], [1, 2, 3]), ('pyarrow', pa.array([1, 2, 3]), [1, 2, 3]), ('pyarrow_ragged', pa.array([[1, 2], [3]]), [1, 2, 3])) def testAccessFeaturesDict(self, feature_value, slice_value): accessor = slice_accessor.SliceAccessor([{'feature': feature_value}]) self.assertEqual(slice_value, accessor.get('feature')) # Test with multiple dicts and duplicate values accessor = slice_accessor.SliceAccessor([{ 'feature': feature_value }, { 'feature': feature_value }]) self.assertEqual(slice_value, accessor.get('feature')) # Test with default features dict accessor = slice_accessor.SliceAccessor( [{ 'unmatched_feature': feature_value }], default_features_dict={'feature': feature_value}) self.assertEqual(slice_value, accessor.get('feature')) def testLegacyAccessFeaturesDict(self): with tf.compat.v1.Session() as sess: sparse = tf.SparseTensor(indices=[[0, 0], [1, 1]], values=['apple', 'banana'], dense_shape=[2, 2]) dense = tf.constant([1.0, 2.0]) dense_single = tf.constant([7.0]) dense_multidim = tf.constant([[1.0, 2.0], [3.0, 4.0]]) squeeze_needed = tf.constant([[2.0]]) (sparse_value, dense_value, dense_single_value, dense_multidim_value, squeeze_needed_value) = sess.run(fetches=[ sparse, dense, dense_single, dense_multidim, squeeze_needed ]) features_dict = { 'sparse': { encoding.NODE_SUFFIX: sparse_value }, 'dense': { encoding.NODE_SUFFIX: dense_value }, 'dense_single': { encoding.NODE_SUFFIX: dense_single_value }, 'squeeze_needed': { encoding.NODE_SUFFIX: squeeze_needed_value }, 'dense_multidim': { encoding.NODE_SUFFIX: dense_multidim_value }, } accessor = slice_accessor.SliceAccessor([features_dict]) self.assertEqual([b'apple', b'banana'], accessor.get('sparse')) self.assertEqual([1.0, 2.0], accessor.get('dense')) self.assertEqual([7.0], accessor.get('dense_single')) self.assertEqual([1.0, 2.0, 3.0, 4.0], accessor.get('dense_multidim')) self.assertEqual([2.0], accessor.get('squeeze_needed'))
def testSplitExtracts(self): extracts = { 'features': { 'feature_1': np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]), 'feature_2': np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]), 'feature_3': types.SparseTensorValue(values=np.array([1, 2, 3]), indices=np.array([[0, 0, 1], [1, 0, 2], [2, 0, 0]]), dense_shape=np.array([3, 1, 3])), 'feature_4': types.RaggedTensorValue(values=np.array([ 3, 1, 4, 1, 5, 9, 2, 6, 3, 1, 4, 1, 5, 9, 2, 6, 3, 1, 4, 1, 5, 9, 2, 6 ]), nested_row_splits=[ np.array([0, 5, 10, 15]), np.array([ 0, 4, 4, 7, 8, 8, 12, 12, 15, 16, 16, 20, 20, 23, 24, 24 ]) ]) }, 'labels': np.array([1.0, 0.0, 1.0]), 'example_weights': np.array([0.0, 0.5, 1.0]), 'predictions': { 'model1': np.array([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]]), 'model2': np.array([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]]) }, '_slice_key_types': np.array([(), (), ()]) } expected = [ { 'features': { 'feature_1': np.array([1.0, 2.0]), 'feature_2': np.array([1.0, 2.0]), 'feature_3': types.SparseTensorValue(values=np.array([1]), indices=np.array([[0, 1]]), dense_shape=np.array([1, 3])), 'feature_4': types.RaggedTensorValue( values=np.array([3, 1, 4, 1, 5, 9, 2, 6]), nested_row_splits=[np.array([0, 4, 4, 7, 8, 8])]) }, 'labels': np.array([1.0]), 'example_weights': np.array([0.0]), 'predictions': { 'model1': np.array([0.1, 0.2]), 'model2': np.array([0.1, 0.2]) }, '_slice_key_types': np.array([()]) }, { 'features': { 'feature_1': np.array([3.0, 4.0]), 'feature_2': np.array([3.0, 4.0]), 'feature_3': types.SparseTensorValue(values=np.array([2]), indices=np.array([[0, 2]]), dense_shape=np.array([1, 3])), 'feature_4': types.RaggedTensorValue( values=np.array([3, 1, 4, 1, 5, 9, 2, 6]), nested_row_splits=[np.array([0, 4, 4, 7, 8, 8])]) }, 'labels': np.array([0.0]), 'example_weights': np.array([0.5]), 'predictions': { 'model1': np.array([0.3, 0.4]), 'model2': np.array([0.3, 0.4]) }, '_slice_key_types': np.array([()]) }, { 'features': { 'feature_1': np.array([5.0, 6.0]), 'feature_2': np.array([5.0, 6.0]), 'feature_3': types.SparseTensorValue(values=np.array([3]), indices=np.array([[0, 0]]), dense_shape=np.array([1, 3])), 'feature_4': types.RaggedTensorValue( values=np.array([3, 1, 4, 1, 5, 9, 2, 6]), nested_row_splits=[np.array([0, 4, 4, 7, 8, 8])]) }, 'labels': np.array([1.0]), 'example_weights': np.array([1.0]), 'predictions': { 'model1': np.array([0.5, 0.6]), 'model2': np.array([0.5, 0.6]) }, '_slice_key_types': np.array([()]) }, ] self.assertAllClose(util.split_extracts(extracts), expected)