def test_array_to_recordio_dense():
    array_data = [[1.0, 2.0, 3.0], [10.0, 20.0, 30.0]]
    buf = _encoders.array_to_recordio_protobuf(np.array(array_data))
    for record_data, expected in zip(_read_recordio(buf), array_data):
        record = Record()
        record.ParseFromString(record_data)
        assert record.features["values"].float64_tensor.values == expected
def test_encode_selected_predictions_protobuf():
    expected_predicted_labels = [[1], [0]]
    expected_probabilities = [[0.4, 0.6], [0.9, 0.1]]

    protobuf_response = serve_utils.encode_selected_predictions(TEST_PREDICTIONS, TEST_KEYS,
                                                                "application/x-recordio-protobuf")
    stream = io.BytesIO(protobuf_response)

    for recordio, predicted_label, probabilities in zip(_read_recordio(stream),
                                                        expected_predicted_labels, expected_probabilities):
        record = Record()
        record.ParseFromString(recordio)
        assert record.label["predicted_label"].float32_tensor.values == predicted_label
        assert all(np.isclose(record.label["probabilities"].float32_tensor.values, probabilities))
def test_sparse_float32_write_spmatrix_to_sparse_tensor():
    n = 4
    array_data = [[1.0, 2.0], [10.0, 30.0], [100.0, 200.0, 300.0, 400.0],
                  [1000.0, 2000.0, 3000.0]]
    keys_data = [[0, 1], [1, 2], [0, 1, 2, 3], [0, 2, 3]]

    flatten_data = list(itertools.chain.from_iterable(array_data))
    y_indices = list(itertools.chain.from_iterable(keys_data))
    x_indices = [[i] * len(keys_data[i]) for i in range(len(keys_data))]
    x_indices = list(itertools.chain.from_iterable(x_indices))

    array = sparse.coo_matrix((flatten_data, (x_indices, y_indices)),
                              dtype="float32")
    buf = _encoders.array_to_recordio_protobuf(array)
    for record_data, expected_data, expected_keys in zip(
            _read_recordio(buf), array_data, keys_data):
        record = Record()
        record.ParseFromString(record_data)
        assert record.features["values"].float32_tensor.values == expected_data
        assert record.features["values"].float32_tensor.keys == expected_keys
        assert record.features["values"].float32_tensor.shape == [n]