def test_reduce(): assert rd('2bd->2d') == (1, ) assert rd('2bd->2') == (1, 2) print('test_reduce: all assertions hold')
from tsalib import reduce_dims as rd def test_reduce(): assert rd('2bd->2d') == (1, ) assert rd('2bd->2') == (1, 2) print('test_reduce: all assertions hold') test_reduce() # In[16]: x: 'btd' = np.random.rand(_B, _T, _D) y = np.mean(x, axis=rd('btd->b')) assert y.shape == (_B, ) # ## Looong warps # In[17]: def warp_long1(): B, T, D, C = get_dim_vars('b t d c') x1: 'btd' = np.ones((B, T, D)) x2: 'btd' = np.ones((B, T, D)) x3: 'btd' = np.ones((B, T, D)) y = warp([x1, x2, x3], '(btd)* -> btdc -> bdtc -> b,d//2,t*2,c', 'jpv') assert y.shape == (B, D // 2, T * 2, C) print('warp_long1: all assertions hold')