def test_batch_empty(): b5_dict = np.array([{'a': False, 'b': {'c': 2.0, 'd': 1.0}}, {'a': True, 'b': {'c': 3.0}}]) b5 = Batch(b5_dict) b5[1] = Batch.empty(b5[0]) assert np.allclose(b5.a, [False, False]) assert np.allclose(b5.b.c, [2, 0]) assert np.allclose(b5.b.d, [1, 0]) data = Batch(a=[False, True], b={'c': np.array([2., 'st'], dtype=object), 'd': [1, None], 'e': [2., float('nan')]}, c=np.array([1, 3, 4], dtype=int), t=torch.tensor([4, 5, 6, 7.])) data[-1] = Batch.empty(data[1]) assert np.allclose(data.c, [1, 3, 0]) assert np.allclose(data.a, [False, False]) assert list(data.b.c) == [2.0, None] assert list(data.b.d) == [1, None] assert np.allclose(data.b.e, [2, 0]) assert torch.allclose(data.t, torch.tensor([4, 5, 6, 0.])) data[0].empty_() # which will fail in a, b.c, b.d, b.e, c assert torch.allclose(data.t, torch.tensor([0., 5, 6, 0])) data.empty_(index=0) assert np.allclose(data.c, [0, 3, 0]) assert list(data.b.c) == [None, None] assert list(data.b.d) == [None, None] assert list(data.b.e) == [0, 0] b0 = Batch() b0.empty_() assert b0.shape == []
def test_batch_cat_and_stack_and_empty(): b1 = Batch(a=[{'b': np.float64(1.0), 'd': Batch(e=np.array(3.0))}]) b2 = Batch(a=[{'b': np.float64(4.0), 'd': {'e': np.array(6.0)}}]) b12_cat_out = Batch.cat((b1, b2)) b12_cat_in = copy.deepcopy(b1) b12_cat_in.cat_(b2) assert np.all(b12_cat_in.a.d.e == b12_cat_out.a.d.e) assert np.all(b12_cat_in.a.d.e == b12_cat_out.a.d.e) assert isinstance(b12_cat_in.a.d.e, np.ndarray) assert b12_cat_in.a.d.e.ndim == 1 b12_stack = Batch.stack((b1, b2)) assert isinstance(b12_stack.a.d.e, np.ndarray) assert b12_stack.a.d.e.ndim == 2 b3 = Batch(a=np.zeros((3, 4)), b=torch.ones((2, 5)), c=Batch(d=[[1], [2]])) b4 = Batch(a=np.ones((3, 4)), b=torch.ones((2, 5)), c=Batch(d=[[0], [3]])) b34_stack = Batch.stack((b3, b4), axis=1) assert np.all(b34_stack.a == np.stack((b3.a, b4.a), axis=1)) assert np.all(b34_stack.c.d == list(map(list, zip(b3.c.d, b4.c.d)))) b5_dict = np.array([{ 'a': False, 'b': { 'c': 2.0, 'd': 1.0 } }, { 'a': True, 'b': { 'c': 3.0 } }]) b5 = Batch(b5_dict) assert b5.a[0] == np.array(False) and b5.a[1] == np.array(True) assert np.all(b5.b.c == np.stack([e['b']['c'] for e in b5_dict], axis=0)) assert b5.b.d[0] == b5_dict[0]['b']['d'] assert b5.b.d[1] == 0.0 b5[1] = Batch.empty(b5[0]) assert np.allclose(b5.a, [False, False]) assert np.allclose(b5.b.c, [2, 0]) assert np.allclose(b5.b.d, [1, 0]) data = Batch(a=[False, True], b={ 'c': [2., 'st'], 'd': [1, None], 'e': [2., float('nan')] }, c=np.array([1, 3, 4], dtype=np.int), t=torch.tensor([4, 5, 6, 7.])) data[-1] = Batch.empty(data[1]) assert np.allclose(data.c, [1, 3, 0]) assert np.allclose(data.a, [False, False]) assert list(data.b.c) == ['2.0', ''] assert list(data.b.d) == [1, None] assert np.allclose(data.b.e, [2, 0]) assert torch.allclose(data.t, torch.tensor([4, 5, 6, 0.])) b0 = Batch() b0.empty_() assert b0.shape == []