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