def test_reversible_permute(self):
layer = sparsity.ReversibleReshapePermute()
x = np.array([[1, 2, 3, 4, 5, 6, 7, 8], [0, 1, 2, 3, 4, 5, 6, 7]])
layer.init(shapes.signature(x))
ys = layer(x)
self.assertEqual(tl.to_list(ys),
[[1, 3, 5, 7, 2, 4, 6, 8], [0, 2, 4, 6, 1, 3, 5, 7]])
rev_x = layer.reverse(ys, weights=layer.weights)
self.assertEqual(tl.to_list(x), tl.to_list(rev_x))
def test_from_file(self):
params = np.array([[0.0, 0.1], [0.2, 0.3], [0.4, 0.5]])
# `create_tempfile` needs access to --test_tmpdir, however in the OSS world
# pytest doesn't run `absltest.main`, so we need to manually parse the flags
test_utils.ensure_flag('test_tmpdir')
filename = self.create_tempfile('params.npy').full_path
with open(filename, 'wb') as f:
np.save(f, params)
f = tl.InitializerFromFile(filename)
init_value = f(params.shape, rng())
self.assertEqual(tl.to_list(init_value), tl.to_list(params))
def test_padding_same(self):
layer = tl.SumPool(pool_size=(3, ), strides=(3, ), padding='SAME')
# One padding position needed; add at end.
x = np.array([[[0, 9], [1, 8], [2, 7], [3, 6], [4, 5]]])
y = layer(x)
self.assertEqual(tl.to_list(y), [[[3, 24], [7, 11]]])
# Two padding positions needed; add one at end and one at start.
x = np.array([[[0, 9], [1, 8], [2, 7], [3, 6]]])
y = layer(x)
self.assertEqual(tl.to_list(y), [[[1, 17], [5, 13]]])
def test_one_hot(self):
targets = np.array([2, 0, 1])
n_categories = 5
target_distributions = tl.one_hot(targets, n_categories)
self.assertEqual(
tl.to_list(target_distributions),
[[0., 0., 1., 0., 0.], [1., 0., 0., 0., 0.], [0., 1., 0., 0., 0.]])
def test_shift_right_cls(self):
layer = ra.ShiftRightCls(5)
x = np.array([[1, 2, 3, 4]])
_, _ = layer.init(shapes.signature(x))
y = layer(x)
self.assertEqual(tl.to_list(y), [[5, 1, 2, 3]])
def test_padding_default(self):
layer = tl.MaxPool(pool_size=(3, ), strides=(3, ))
# Discard incomplete window at end: [[3, 6], [4, 5]].
x = np.array([[[0, 9], [1, 8], [2, 7], [3, 6], [4, 5]]])
y = layer(x)
self.assertEqual(tl.to_list(y), [[[2, 9]]])
def test_create_mask_layer(self):
layer = ra.CreateAttentionMaskLayer()
xs = _get_xs(q=2, k=2)
layer.init(shapes.signature(xs))
_, _, _, mask = layer(xs)
self.assertEqual(mask.shape, (1, 1, 2, 2))
np.testing.assert_equal(tl.to_list(mask),
[[[[True, False], [True, True]]]])
def test_create_mask_layer(self):
layer = ra.AttentionMaskLayer()
xs = np.zeros((1, 2, 5))
layer.init(shapes.signature(xs))
mask = layer(xs)
self.assertEqual(mask.shape, (2, 2))
np.testing.assert_equal(tl.to_list(mask),
[[True, False], [True, True]])
def test_forward(self):
layer = tl.AvgPool(pool_size=(2, 2), strides=(2, 2))
x = np.array([[
[[1, 2, 3], [4, 5, 6], [10, 20, 30], [40, 50, 60]],
[[4, 2, 3], [7, 1, 2], [40, 20, 30], [70, 10, 20]],
]])
y = layer(x)
self.assertEqual(tl.to_list(y), [[[[4.0, 2.5, 3.5], [40, 25, 35]]]])
def test_fast_shift_matrix_funnel_2(self):
layer = ra._fast_matrix_shift
x = np.array([[[[-3., -2., -1., 0., 1., 2., 3.],
[-3., -2., -1., 0., 1., 2., 3.]]]]).astype(np.float32)
y = layer(x, funnel_factor=2, is_upsampling=False)
self.assertEqual(y.dtype, np.float32)
self.assertEqual(tl.to_list(y),
[[[[0., 1., 2., 3.], [-2., -1., 0., 1.]]]])
def test_fast_shift_matrix(self):
layer = ra._fast_matrix_shift
x = np.array([[[[-3., -2., -1., 0.], [-3., -2., -1., 0.],
[-3., -2., -1., 0.], [-3., -2., -1.,
0.]]]]).astype(np.float32)
y = layer(x)
self.assertEqual(y.dtype, np.float32)
self.assertEqual(tl.to_list(y),
[[[[0., 0., -3., -2.], [-1., 0., 0., -3.],
[-2., -1., 0., 0.], [-3., -2., -1., 0.]]]])
def test_shift_right(self):
# Test shifts right on axis=1
layer = tl.ShiftRight()
x = np.array([[[9, 9, 9], [8, 8, 8], [7, 7, 7], [6, 6, 6]],
[[99, 98, 97], [96, 95, 94], [93, 92, 91], [90, 89,
88]]])
y = layer(x)
self.assertEqual(x.shape, y.shape)
self.assertEqual(
tl.to_list(y),
[[[0, 0, 0], [9, 9, 9], [8, 8, 8], [7, 7, 7]],
[[0, 0, 0], [99, 98, 97], [96, 95, 94], [93, 92, 91]]])
def test_create_mask_layer_predict(self):
layer = ra.AttentionMaskLayer(total_kv_pooling=2,
n_raw_tokens_generated=1,
max_inference_length=3,
mode='predict')
xs = np.zeros((1, 1, 5))
layer.init(shapes.signature(xs))
for _ in range(2):
mask = layer(xs)
self.assertEqual(mask.shape, (1, 3))
np.testing.assert_equal(tl.to_list(mask), [[True, False, False]])
for _ in range(2):
mask = layer(xs)
self.assertEqual(mask.shape, (1, 3))
np.testing.assert_equal(tl.to_list(mask), [[True, True, False]])
for _ in range(2):
mask = layer(xs)
self.assertEqual(mask.shape, (1, 3))
np.testing.assert_equal(tl.to_list(mask), [[True, True, True]])
def test_reversible_permute(self):
layer = sparsity.ReversibleRandomPermute()
x = np.array([[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13],
[0, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0, 11, 12, 13],
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13],
])
layer.init(shapes.signature(x))
ys = layer(x)
# this assert will fail once per ~87B runs, but it's okay
self.assertNotEqual(tl.to_list(ys), tl.to_list(x))
self.assertEqual(tl.to_list(ys[0]), tl.to_list(ys[2]))
self.assertNotEqual(tl.to_list(ys[0]), tl.to_list(ys[1]))
rev_x = layer.reverse(ys, weights=layer.weights)
self.assertEqual(tl.to_list(x), tl.to_list(rev_x))
def test_shift_right_float(self):
layer = tl.ShiftRight()
x = np.array([[[9, 9, 9], [8, 8, 8], [7, 7, 7], [6, 6, 6]],
[[99, 98, 97], [96, 95, 94], [93, 92, 91],
[90, 89, 88]]]).astype(np.float32)
x /= 2.0
self.assertEqual(x.dtype, np.float32)
y = layer(x)
self.assertEqual(y.dtype, np.float32)
self.assertEqual(tl.to_list(y),
[[[0.0, 0.0, 0.0], [4.5, 4.5, 4.5], [4.0, 4.0, 4.0],
[3.5, 3.5, 3.5]],
[[0.0, 0.0, 0.0], [49.5, 49.0, 48.5],
[48.0, 47.5, 47.0], [46.5, 46.0, 45.5]]])
def test_reversible_swap(self):
layer = tl.ReversibleSwap()
xs = [np.array([1, 2]), np.array([10, 20])]
ys = layer(xs)
self.assertEqual(tl.to_list(ys), [[10, 20], [1, 2]])
def test_reversible_swap(self, backend):
with fastmath.use_backend(backend):
layer = tl.ReversibleSwap()
xs = [np.array([1, 2]), np.array([10, 20])]
ys = layer(xs)
self.assertEqual(tl.to_list(ys), [[10, 20], [1, 2]])
def test_relu(self):
layer = tl.Relu()
x = np.array([-2.0, -1.0, 0.0, 2.0, 3.0, 5.0])
y = layer(x)
self.assertEqual(tl.to_list(y), [0.0, 0.0, 0.0, 2.0, 3.0, 5.0])
def test_parametric_relu(self):
layer = tl.ParametricRelu(a=.25)
x = np.array([-2.0, -1.0, 0.0, 2.0, 3.0, 5.0])
y = layer(x)
self.assertEqual(tl.to_list(y), [0.0, 0.0, 0.0, .5, .75, 1.25])
def test_hard_tanh(self):
layer = tl.HardTanh()
x = np.array([-1.5, -.5, -.25, 0.0, .25, .5, 1.5])
y = layer(x)
self.assertEqual(tl.to_list(y), [-1.0, -.5, -.25, 0.0, .25, .5, 1.0])
def test_hard_sigmoid(self):
layer = tl.HardSigmoid()
x = np.array([-1.5, -.5, -.25, 0.0, .25, .5, 1.5])
y = layer(x)
self.assertEqual(tl.to_list(y), [0.0, 0.5, 0.75, 1.0, 1.0, 1.0, 1.0])
def test_leaky_relu(self):
layer = tl.LeakyRelu(a=.125)
x = np.array([-2.0, -1.0, 0.0, 2.0, 3.0, 5.0])
y = layer(x)
self.assertEqual(tl.to_list(y), [-.25, -.125, 0.0, 2.0, 3.0, 5.0])