def test_pad_arrays_axis(axis: int, is_right_pad: bool): arrays = [d["target"] for d in list(iter(get_dataset()))] if axis == 0: arrays = [x.T for x in arrays] padded_arrays = _pad_arrays(arrays, axis, is_right_pad=is_right_pad) assert all(a.shape[axis] == 8 for a in padded_arrays) assert all(a.shape[1 - axis] == 2 for a in padded_arrays)
def test_variable_length_stack(is_right_pad): arrays = [d["target"].T for d in list(iter(get_dataset()))] stacked = stack( arrays, variable_length=True, is_right_pad=is_right_pad, ) assert stacked.shape[0] == 3 assert stacked.shape[1] > 0 assert stacked.shape[2] == 2
def test_pad_arrays_axis(array_type, multi_processing, axis: int): arrays = [ d["target"] if array_type == "np" else mx.nd.array(d["target"]) for d in list(iter(get_dataset())) ] if axis == 0: arrays = [x.T for x in arrays] padded_arrays = _pad_arrays(arrays, axis) assert all(a.shape[axis] == 8 for a in padded_arrays) assert all(a.shape[1 - axis] == 2 for a in padded_arrays)
def test_inference_loader_short_intervals(loader_factory): loader = loader_factory( dataset=get_dataset(), prediction_interval_length=0.001, context_interval_length=0.0001, is_train=False, override_args={"batch_size": 5}, ) batches = list(loader) batch = batches[0] assert batch["past_target"].shape[1] == 1
def test_inference_loader_shapes_small_batch(loader_factory): loader = loader_factory( dataset=get_dataset(), prediction_interval_length=1.0, context_interval_length=1.5, is_train=False, override_args={"batch_size": 2}, ) batches = list(loader) assert len(batches) == 2 batch = batches[0] assert batch["past_target"].shape[2] == 2 assert batch["past_target"].shape[0] == 2 assert batch["past_valid_length"].shape[0] == 2
def test_train_loader_shapes(loader_factory): loader = loader_factory(get_dataset(), 1.0, 1.5) d = next(iter(loader)) field_names = [ "past_target", "past_valid_length", "future_target", "future_valid_length", ] assert all([key in d for key in field_names]) assert d["past_target"].shape[2] == d["future_target"].shape[2] == 2 assert d["past_target"].shape[0] == d["future_target"].shape[0] == 10 assert (d["past_valid_length"].shape[0] == d["future_valid_length"].shape[0] == 10)
def test_pad_arrays_pad_left(): arrays = [d["target"] for d in list(iter(get_dataset()))] padded_arrays = _pad_arrays(arrays, 1, is_right_pad=False) for padded_array in padded_arrays[1:]: assert np.allclose(padded_array[:, 0], 0)
def test_train_loader_length(loader_factory): loader = loader_factory(get_dataset(), 1.0, 1.5) assert len(list(loader)) == NUM_BATCHES