def test_shuffle_layer(self): shuffle_layer = rse.ShuffleLayer() x = np.array([[[0], [1], [2], [3], [4], [5], [6], [7]]]) print(x.shape) _, _ = shuffle_layer.init(shapes.signature(x)) y = shuffle_layer(x) expected_output = np.array([[[0], [2], [4], [6], [1], [3], [5], [7]]]) self._assert_equal_tensors(y, expected_output)
def test_shuffle_layer_log_times_is_identity(self):
seq_len = 8
d_model = 17
shuffle_layer = rse.ShuffleLayer()
x = _input_with_indice_as_values(seq_len, d_model)
_, _ = shuffle_layer.init(shapes.signature(x))
y = x
for _ in range(np.int(np.log2(seq_len))):
y = shuffle_layer(y)
self._assert_equal_tensors(x, y)