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)
