def test_rows(self): s1 = SampleBatch( { "a": np.array([[1, 1], [2, 2], [3, 3]]), "b": {"c": np.array([[4, 4], [5, 5], [6, 6]])}, SampleBatch.SEQ_LENS: np.array([1, 2]), } ) check( next(s1.rows()), {"a": [1, 1], "b": {"c": [4, 4]}, SampleBatch.SEQ_LENS: 1}, )
def test_compression(self): """Tests, whether compression and decompression work properly.""" s1 = SampleBatch({ "a": np.array([1, 2, 3, 2, 3, 4]), "b": { "c": np.array([4, 5, 6, 5, 6, 7]) }, }) # Test, whether compressing happens in-place. s1.compress(columns={"a", "b"}, bulk=True) self.assertTrue(is_compressed(s1["a"])) self.assertTrue(is_compressed(s1["b"]["c"])) self.assertTrue(isinstance(s1["b"], dict)) # Test, whether de-compressing happens in-place. s1.decompress_if_needed(columns={"a", "b"}) check(s1["a"], [1, 2, 3, 2, 3, 4]) check(s1["b"]["c"], [4, 5, 6, 5, 6, 7]) it = s1.rows() next(it) check(next(it), {"a": 2, "b": {"c": 5}})