def test_get_item(data):
dlen = get_dataset_lens(data)[0]
assert dlen == 3
assert len(set(get_dataset_lens(data))) == 1
assert get_dataset_item(data, 1) == {
"a": [1],
"b": {
"d": 1
},
"c": np.array([1])
}
def nested_numpy_minibatch(data, batch_size=1):
lens = get_dataset_lens(data)
if isinstance(lens, collections.Mapping):
ln = [v for v in lens.values()][0]
elif isinstance(lens, collections.Sequence):
ln = lens[0]
else:
ln = lens
for idx in BatchSampler(range(ln),
batch_size=batch_size,
drop_last=False):
yield get_dataset_item(data, idx)
def test_datset_lens_good(data):
assert get_dataset_lens(data) == 3 * [3]
def test_datset_lens_bad(bad_data):
assert sorted(get_dataset_lens(bad_data)) == [1, 3, 3, 4]
with pytest.raises(Exception):
get_dataset_lens(bad_data, require_numpy=True)
def __init__(self, *args, **kwargs):
self.data = self._get_data_fn()(*args, **kwargs)
lens = get_dataset_lens(self.data, require_numpy=True)
# check that all dimensions are the same
assert len(set(lens)) == 1
self.n = lens[0]