# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import math import random import numpy as np import pytest import torch from nemo.collections.tts.modules.transformer_2501 import ( ConvolutionLayer, CrossAttention, PositionwiseConvFF, SelfAttention, Transformer, TransformerLayer, ) from nemo.collections.tts.parts.utils.tts_dataset_utils import beta_binomial_prior_distribution def set_seed(seed): torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) np.random.seed(seed) random.seed(seed) @pytest.mark.unit class TestConvolutionLayer: @classmethod def setup_class(cls): cls.in_channels = 3 cls.out_channels = 6 cls.kernel_size = 3 cls.stride = 1 cls.dilation = 1 cls.bias = True # fmt:off cls.input_tensor = torch.Tensor( [[[-1.0542, 0.2675, 0.6963, 0.4738, 0.3910, -0.1505, 0.9171, -0.1528, 3.7269, 0.1779], [-1.0317, 1.6818, 1.4257, -0.5003, -1.7254, 0.8830, -0.4541, -0.4631, -0.0986, 0.5083], [-0.3231, -1.0899, 0.5774, 0.1661, 0.9620, -2.3307, -0.6158, -0.3663, 1.2469, -1.0208]]] ) # fmt:on def test_non_causal_forward(self): set_seed(0) layer = ConvolutionLayer( self.in_channels, self.out_channels, self.kernel_size, self.stride, dilation=self.dilation, bias=self.bias, is_causal=False, ) with torch.no_grad(): output_tensor = layer(self.input_tensor) # fmt:off expected_output_tensor = torch.Tensor( [[[ 0.1912, -0.0555, -0.2681, -0.2289, 1.0788, -0.3908, 0.0936, -0.7962, 1.3754, -0.0731], [-0.3715, -0.6326, -0.9596, -0.0933, -0.1024, -0.2082, -0.5924, 0.1097, -0.5418, -0.0854], [ 0.3974, 0.4537, 0.3299, 0.1471, -0.5983, -0.8645, 0.0975, 0.6063, -0.6619, -0.9711], [-0.3048, 0.3862, -0.2462, -0.9903, -0.6189, 0.7389, 0.0785, -1.0870, -1.0018, -1.2426], [-0.4357, -0.0446, 0.0879, 0.0930, -0.2242, 0.5285, 0.4006, -0.1846, 0.5668, -0.5242], [-0.0625, 0.4123, -0.6289, -0.4317, 0.1595, 0.0386, -1.0774, 0.2218, 0.8483, -0.4886]]] ) # fmt:on assert torch.allclose(output_tensor, expected_output_tensor, atol=1e-4) def test_causal_forward(self): set_seed(0) layer = ConvolutionLayer( self.in_channels, self.out_channels, self.kernel_size, self.stride, dilation=self.dilation, bias=self.bias, is_causal=True, ) with torch.no_grad(): output_tensor = layer(self.input_tensor) # fmt:off expected_output_tensor = torch.Tensor( [[[ 0.4301, 0.1912, -0.0555, -0.2681, -0.2289, 1.0788, -0.3908, 0.0936, -0.7962, 1.3754], [-0.0501, -0.3715, -0.6326, -0.9596, -0.0933, -0.1024, -0.2082, -0.5924, 0.1097, -0.5418], [-0.4204, 0.3974, 0.4537, 0.3299, 0.1471, -0.5983, -0.8645, 0.0975, 0.6063, -0.6619], [ 0.1543, -0.3048, 0.3862, -0.2462, -0.9903, -0.6189, 0.7389, 0.0785, -1.0870, -1.0018], [-0.1337, -0.4357, -0.0446, 0.0879, 0.0930, -0.2242, 0.5285, 0.4006, -0.1846, 0.5668], [-0.6127, -0.0625, 0.4123, -0.6289, -0.4317, 0.1595, 0.0386, -1.0774, 0.2218, 0.8483]]] ) # fmt:on assert torch.allclose(output_tensor, expected_output_tensor, atol=1e-4) class TestPositionwiseConvFF: @classmethod def setup_class(cls): cls.d_model = 3 cls.d_ffn = 12 cls.p_dropout = 0.0 cls.kernel_size = 3 cls.bias = True # fmt:off cls.input_tensor = torch.Tensor( [[[-1.6682, -0.6069, 0.1321], [-1.5489, 0.3279, -0.9159], [-0.7490, 1.8984, 0.5030], [-0.8130, 0.0058, -1.9979], [-1.4994, -0.3270, 1.4961], [-1.6613, -1.7827, 0.8932], [-0.6276, -1.0770, -0.9971], [ 1.5424, 1.3590, 1.2287], [-0.1543, 0.3365, 1.7475], [-0.1753, 0.4115, 0.0772]]] ) # fmt:on def test_causal_forward(self): set_seed(0) layer = PositionwiseConvFF( self.d_model, self.d_ffn, self.p_dropout, self.kernel_size, bias=self.bias, is_causal=True ) with torch.no_grad(): output_tensor = layer(self.input_tensor) # fmt:off expected_output_tensor = torch.Tensor( [[[-0.1242, -0.0114, 0.0212], [-0.0441, -0.0555, -0.0795], [-0.0282, 0.0366, -0.2033], [-0.0421, 0.0305, -0.2573], [-0.1877, -0.2492, -0.1638], [-0.4300, -0.1160, 0.2177], [-0.1652, -0.3130, -0.3329], [-0.1737, 0.1133, -0.1802], [-0.2599, -0.0381, 0.1362], [-0.0584, -0.2936, 0.2719]]] ) # fmt:on assert torch.allclose(output_tensor, expected_output_tensor, atol=1e-4) def test_non_causal_forward(self): set_seed(0) layer = PositionwiseConvFF( self.d_model, self.d_ffn, self.p_dropout, self.kernel_size, bias=self.bias, is_causal=False ) with torch.no_grad(): output_tensor = layer(self.input_tensor) # fmt:off expected_output_tensor = torch.Tensor( [[[-0.0617, -0.0321, -0.1646], [-0.0421, 0.0305, -0.2573], [-0.1877, -0.2492, -0.1638], [-0.4300, -0.1160, 0.2177], [-0.1652, -0.3130, -0.3329], [-0.1737, 0.1133, -0.1802], [-0.2599, -0.0381, 0.1362], [-0.0584, -0.2936, 0.2719], [ 0.0361, 0.1110, 0.0441], [-0.0244, 0.0682, 0.0340]]] ) # fmt:on assert torch.allclose(output_tensor, expected_output_tensor, atol=1e-4) class TestSelfAttention: @classmethod def setup_class(cls): cls.n_heads = 2 cls.d_model = 4 cls.p_dropout = 0.0 cls.max_length_causal_mask = 6 # fmt:off cls.query_tensor = torch.Tensor( [[[ 0.7239, -0.2362, -0.6610, -1.3759], [ 1.7381, 0.0793, -1.1241, 0.9529], [-1.9809, 0.2217, 0.0795, 0.0307], [ 0.3208, 0.4485, 0.3046, -0.0704], [-1.4412, 0.8981, 0.1219, 0.0481], [ 1.7811, -0.1358, 0.6073, 0.8275]]] ) # fmt:on def test_causal_forward(self): set_seed(0) layer = SelfAttention( self.n_heads, self.d_model, self.p_dropout, is_causal=True, max_length_causal_mask=self.max_length_causal_mask, ) query_mask = torch.ones(1, self.max_length_causal_mask).bool() with torch.no_grad(): output_tensor, attn_output = layer(self.query_tensor, query_mask) # fmt:off expected_output_tensor = torch.Tensor( [[[-0.2569, 0.2782, -0.0348, 0.2480], [-0.3949, 0.4054, -0.0876, 0.2574], [-0.1033, 0.0659, -0.0259, 0.1738], [-0.1485, 0.0995, -0.0415, 0.0684], [-0.0123, 0.0185, -0.0027, 0.0708], [-0.0672, 0.0566, -0.0214, -0.0021]]] ) expected_attn_prob = torch.Tensor( [[[[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000], [0.5828, 0.4172, 0.0000, 0.0000, 0.0000, 0.0000], [0.3216, 0.4260, 0.2525, 0.0000, 0.0000, 0.0000], [0.2385, 0.2238, 0.2872, 0.2504, 0.0000, 0.0000], [0.1807, 0.1973, 0.2057, 0.2045, 0.2118, 0.0000], [0.1159, 0.1388, 0.2010, 0.1721, 0.2161, 0.1562]], [[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000], [0.2866, 0.7134, 0.0000, 0.0000, 0.0000, 0.0000], [0.2799, 0.2472, 0.4729, 0.0000, 0.0000, 0.0000], [0.2964, 0.2535, 0.2075, 0.2427, 0.0000, 0.0000], [0.1864, 0.1616, 0.2394, 0.1974, 0.2152, 0.0000], [0.1666, 0.2030, 0.1391, 0.1649, 0.1546, 0.1719]]]] ) expected_attn_score = torch.Tensor( [[[[ 0.5248, float('-inf'), float('-inf'), float('-inf'), float('-inf'), float('-inf')], [-0.1948, -0.5291, float('-inf'), float('-inf'), float('-inf'), float('-inf')], [-0.0279, 0.2533, -0.2698, float('-inf'), float('-inf'), float('-inf')], [-0.0508, -0.1145, 0.1350, -0.0020, float('-inf'), float('-inf')], [-0.0985, -0.0105, 0.0315, 0.0257, 0.0604, float('-inf')], [-0.3253, -0.1457, 0.2250, 0.0694, 0.2971, -0.0275]], [[ 0.5075, float('-inf'), float('-inf'), float('-inf'), float('-inf'), float('-inf')], [-0.5541, 0.3578, float('-inf'), float('-inf'), float('-inf'), float('-inf')], [-0.2499, -0.3738, 0.2746, float('-inf'), float('-inf'), float('-inf')], [ 0.2215, 0.0654, -0.1351, 0.0216, float('-inf'), float('-inf')], [-0.1011, -0.2439, 0.1488, -0.0438, 0.0425, float('-inf')], [ 0.0526, 0.2502, -0.1277, 0.0424, -0.0221, 0.0840]]]] ) # fmt:on assert torch.allclose(output_tensor, expected_output_tensor, atol=1e-4) assert torch.allclose(attn_output[0], expected_attn_prob, atol=1e-4) assert torch.allclose(attn_output[1], expected_attn_score, atol=1e-4) def test_non_causal_forward(self): set_seed(0) layer = SelfAttention( self.n_heads, self.d_model, self.p_dropout, is_causal=False, max_length_causal_mask=self.max_length_causal_mask, ) query_mask = torch.ones(1, self.max_length_causal_mask).bool() with torch.no_grad(): output_tensor, attn_output = layer(self.query_tensor, query_mask) # fmt:off expected_output_tensor = torch.Tensor( [[[-0.0954, 0.1131, -0.0195, 0.0704], [-0.0401, 0.0364, -0.0156, 0.0088], [ 0.0324, 0.0368, 0.0174, 0.0501], [-0.0633, 0.0610, -0.0176, 0.0180], [-0.0017, 0.0361, 0.0030, 0.0319], [-0.0672, 0.0566, -0.0214, -0.0021]]] ) expected_attn_prob = torch.Tensor( [[[[0.2835, 0.1482, 0.1723, 0.1426, 0.1422, 0.1111], [0.1254, 0.0898, 0.2819, 0.1493, 0.2682, 0.0854], [0.1549, 0.2051, 0.1216, 0.1663, 0.1294, 0.2228], [0.1583, 0.1485, 0.1906, 0.1662, 0.1890, 0.1475], [0.1498, 0.1635, 0.1705, 0.1696, 0.1755, 0.1711], [0.1159, 0.1388, 0.2010, 0.1721, 0.2161, 0.1562]], [[0.2536, 0.1945, 0.1079, 0.1628, 0.1233, 0.1579], [0.0866, 0.2156, 0.1709, 0.1551, 0.1903, 0.1815], [0.1361, 0.1202, 0.2300, 0.1626, 0.1941, 0.1569], [0.2038, 0.1744, 0.1427, 0.1669, 0.1488, 0.1634], [0.1565, 0.1357, 0.2010, 0.1657, 0.1807, 0.1604], [0.1666, 0.2030, 0.1391, 0.1649, 0.1546, 0.1719]]]] ) expected_attn_score = torch.Tensor( [[[[ 5.2482e-01, -1.2346e-01, 2.7022e-02, -1.6210e-01, -1.6488e-01, -4.1190e-01], [-1.9484e-01, -5.2910e-01, 6.1538e-01, -2.0263e-02, 5.6540e-01, -5.7875e-01], [-2.7873e-02, 2.5326e-01, -2.6980e-01, 4.3127e-02, -2.0759e-01, 3.3584e-01], [-5.0756e-02, -1.1455e-01, 1.3498e-01, -2.0208e-03, 1.2687e-01, -1.2113e-01], [-9.8482e-02, -1.0463e-02, 3.1513e-02, 2.5712e-02, 6.0431e-02, 3.4494e-02], [-3.2530e-01, -1.4574e-01, 2.2503e-01, 6.9375e-02, 2.9711e-01, -2.7542e-02]], [[ 5.0748e-01, 2.4198e-01, -3.4704e-01, 6.4137e-02, -2.1347e-01, 3.3762e-02], [-5.5410e-01, 3.5778e-01, 1.2559e-01, 2.8689e-02, 2.3316e-01, 1.8558e-01], [-2.4989e-01, -3.7383e-01, 2.7462e-01, -7.2002e-02, 1.0508e-01, -1.0771e-01], [ 2.2154e-01, 6.5375e-02, -1.3510e-01, 2.1609e-02, -9.3194e-02, 3.4042e-04], [-1.0105e-01, -2.4395e-01, 1.4884e-01, -4.3842e-02, 4.2481e-02, -7.6735e-02], [ 5.2595e-02, 2.5018e-01, -1.2765e-01, 4.2375e-02, -2.2093e-02, 8.4005e-02]]]] ) # fmt:on assert torch.allclose(output_tensor, expected_output_tensor, atol=1e-4) assert torch.allclose(attn_output[0], expected_attn_prob, atol=1e-4) assert torch.allclose(attn_output[1], expected_attn_score, atol=1e-5) class TestCrossAttention: @classmethod def setup_class(cls): cls.n_heads = 2 cls.d_model = 4 cls.d_memory = 3 cls.p_dropout = 0.0 cls.max_length = 6 # fmt:off # shape = (1, cls.max_length, cls.d_model) cls.query_tensor = torch.Tensor( [[[0.7352, -0.5871, -0.1204, -2.0200], [-0.4618, -0.3604, 1.2287, -0.3434], [0.7838, -0.7646, -1.3349, -0.1538], [-0.9749, -1.0789, -0.0126, -0.7225], [2.6929, -0.2091, 2.1242, -1.0123], [0.5094, -2.0566, 1.3922, -0.2156]]] ) # fmt:on cls.query_mask = torch.ones(1, cls.query_tensor.shape[1]).bool() # fmt:off # shape = (1, 5, cls.d_memory) cls.memory_tensor = torch.Tensor( [[[ 2.0132e-01, -5.6582e-01, 1.1191e+00], [-6.2371e-01, -9.3398e-02, -1.3744e+00], [-9.8265e-01, -8.1742e-01, 4.5611e-01], [-5.4802e-01, -1.1218e+00, 7.6138e-01], [-1.9899e+00, -1.7910e-03, 9.0718e-01]]] ) # fmt:on cls.memory_mask = torch.ones(1, cls.memory_tensor.shape[1]).bool() # shape = (1, cls.query_tensor.shape[1], cls.memory_tensor.shape[1]) cls.attn_prior = torch.from_numpy( beta_binomial_prior_distribution( phoneme_count=cls.memory_tensor.shape[1], mel_count=cls.query_tensor.shape[1] ) ).unsqueeze(0) def test_forward(self): set_seed(0) layer = CrossAttention(self.n_heads, self.d_model, self.d_memory, self.p_dropout) with torch.no_grad(): output_tensor, attn_output = layer( self.query_tensor, self.query_mask, self.memory_tensor, self.memory_mask, self.attn_prior ) # fmt:off expected_output_tensor = torch.Tensor( [[[ 0.2267, -0.2271, 0.0573, -0.0681], [ 0.2672, -0.1823, 0.0722, -0.0859], [ 0.3212, -0.2218, 0.0835, -0.0715], [ 0.3568, -0.2573, 0.0918, -0.0789], [ 0.3962, -0.4112, 0.0816, -0.1972], [ 0.3457, -0.4253, 0.0568, -0.2216]]] ) expected_attn_prob = torch.Tensor( [[[[0.4220, 0.4859, 0.0709, 0.0188, 0.0025], [0.3944, 0.3475, 0.1642, 0.0784, 0.0155], [0.1335, 0.3448, 0.2794, 0.1752, 0.0671], [0.0914, 0.3300, 0.2343, 0.2437, 0.1006], [0.0256, 0.1138, 0.2145, 0.3343, 0.3119], [0.0117, 0.0617, 0.1112, 0.3354, 0.4800]], [[0.8045, 0.1024, 0.0661, 0.0242, 0.0028], [0.4020, 0.2953, 0.1914, 0.0907, 0.0207], [0.1446, 0.2798, 0.3026, 0.1965, 0.0766], [0.0718, 0.2151, 0.2778, 0.2719, 0.1634], [0.0673, 0.0341, 0.1929, 0.4534, 0.2522], [0.0064, 0.0264, 0.0999, 0.2872, 0.5802]]]] ) expected_attn_score = torch.Tensor( [[[[-0.5044, 0.4476, -0.4961, -0.5728, -0.8010], [-0.0761, -0.2027, -0.5103, -0.4385, -0.6724], [-0.1525, 0.2576, 0.0471, -0.0138, 0.0075], [-0.3084, -0.0058, -0.7538, -0.7144, -1.0604], [-0.0799, 0.0260, -0.1511, -0.1493, -0.2187], [-0.1935, -0.3225, -0.9867, -0.8637, -1.3162]], [[ 0.4704, -0.7801, -0.2374, 0.0126, -0.3430], [ 0.0623, -0.2461, -0.2380, -0.1743, -0.2658], [ 0.0104, 0.1313, 0.2096, 0.1834, 0.2217], [-0.0859, 0.0300, -0.1194, -0.1410, -0.1111], [ 0.7876, -1.2782, -0.3570, 0.0556, -0.5312], [ 0.1046, -0.2652, -0.1856, -0.1104, -0.2180]]]] ) # fmt:on assert torch.allclose(output_tensor, expected_output_tensor, atol=1e-4) assert torch.allclose(attn_output[0], expected_attn_prob, atol=1e-4) assert torch.allclose(attn_output[1], expected_attn_score, atol=1e-4) class TestTransformerLayer: @classmethod def setup_class(cls): cls.d_model = 2 cls.d_ffn = 8 cls.sa_n_heads = 2 cls.kernel_size = 3 cls.p_dropout = 0.0 cls.max_length_causal_mask = 5 # fmt:off # shape = (1, cls.max_length_causal_mask, cls.d_model) cls.x = torch.Tensor( [[[ 0.5115, 0.0889], [-0.8568, -2.9632], [-1.3728, 0.7325], [-2.4593, -0.9018], [ 0.9621, 0.4212]]] ) # fmt:on cls.x_mask = torch.ones(1, cls.max_length_causal_mask).bool() # fmt:off # shape = (1, 3, cls.d_model) cls.cond = torch.Tensor( [[[ 1.4441, 0.1393], [ 0.2828, -0.2456], [-0.3075, 0.6581]]] ) # fmt:on cls.cond_mask = torch.ones(1, cls.cond.shape[1]).bool() # shape = (1, cls.x.shape[1], cls.cond.shape[1]) cls.attn_prior = torch.from_numpy( beta_binomial_prior_distribution(phoneme_count=cls.cond.shape[1], mel_count=cls.x.shape[1]) ).unsqueeze(0) def test_forward_causal_self_attn_and_has_xattn(self): set_seed(0) layer = TransformerLayer( self.d_model, self.d_ffn, self.sa_n_heads, self.kernel_size, self.p_dropout, has_xattn=True, xa_n_heads=2, xa_d_memory=2, is_causal=True, max_length_causal_mask=self.max_length_causal_mask, ) with torch.no_grad(): output_dict = layer(self.x, self.x_mask, self.cond, self.cond_mask, self.attn_prior) # fmt:off expected_output = { 'output': torch.Tensor( [[[ 0.1936, 0.5387], [-1.0270, -2.5452], [-1.6884, 0.8765], [-2.7496, -0.7887], [ 1.2837, 0.1172]]] ), 'attn_probabilities': { 'self_attn_probabilities': [ torch.Tensor( [[[[1.0000, 0.0000, 0.0000, 0.0000, 0.0000], [0.5000, 0.5000, 0.0000, 0.0000, 0.0000], [0.3068, 0.3068, 0.3864, 0.0000, 0.0000], [0.2213, 0.2213, 0.2787, 0.2787, 0.0000], [0.2180, 0.2180, 0.1730, 0.1730, 0.2180]], [[1.0000, 0.0000, 0.0000, 0.0000, 0.0000], [0.5000, 0.5000, 0.0000, 0.0000, 0.0000], [0.3237, 0.3237, 0.3527, 0.0000, 0.0000], [0.2393, 0.2393, 0.2607, 0.2607, 0.0000], [0.2068, 0.2068, 0.1898, 0.1898, 0.2068]]]] ), torch.Tensor( [[[[0.1154, float('-inf'), float('-inf'), float('-inf'), float('-inf')], [0.1154, 0.1154, float('-inf'), float('-inf'), float('-inf')], [-0.1154, -0.1154, 0.1154, float('-inf'), float('-inf')], [-0.1154, -0.1154, 0.1154, 0.1154, float('-inf')], [0.1154, 0.1154, -0.1154, -0.1154, 0.1154]], [[0.0429, float('-inf'), float('-inf'), float('-inf'), float('-inf')], [0.0429, 0.0429, float('-inf'), float('-inf'), float('-inf')], [-0.0429, -0.0429, 0.0429, float('-inf'), float('-inf')], [-0.0429, -0.0429, 0.0429, 0.0429, float('-inf')], [0.0429, 0.0429, -0.0429, -0.0429, 0.0429]]]] ) ], 'cross_attn_probabilities': [ torch.Tensor( [[[[0.7181, 0.2394, 0.0426], [0.4843, 0.3874, 0.1283], [0.2753, 0.4129, 0.3118], [0.1344, 0.3583, 0.5074], [0.0520, 0.2599, 0.6882]], [[0.5959, 0.1987, 0.2054], [0.2837, 0.2270, 0.4893], [0.3740, 0.5610, 0.0651], [0.2355, 0.6280, 0.1365], [0.0108, 0.0542, 0.9349]]]] ), torch.Tensor( [[[[0.0586, 0.0586, -0.0586], [0.0624, 0.0624, -0.0624], [-0.0624, -0.0624, 0.0624], [-0.0624, -0.0624, 0.0624], [0.0624, 0.0624, -0.0624]], [[-0.8214, -0.8214, 0.8214], [-0.8745, -0.8744, 0.8745], [0.8745, 0.8744, -0.8745], [0.8745, 0.8744, -0.8745], [-0.8744, -0.8744, 0.8744]]]] ) ] } } # fmt:on assert torch.allclose(output_dict["output"], expected_output["output"], atol=1e-4) for i in range(2): assert torch.allclose( output_dict["attn_probabilities"]["self_attn_probabilities"][i], expected_output["attn_probabilities"]["self_attn_probabilities"][i], atol=1e-4, ) assert torch.allclose( output_dict["attn_probabilities"]["cross_attn_probabilities"][i], expected_output["attn_probabilities"]["cross_attn_probabilities"][i], atol=1e-4, ) @pytest.mark.unit class TestTransformer: @classmethod def setup_class(cls): cls.n_layers = 1 cls.d_model = 4 cls.d_ffn = 16 cls.sa_n_heads = 2 cls.kernel_size = 3 cls.p_dropout = 0.0 cls.p_dropout_out = 0.0 cls.is_causal = True cls.max_length_causal_mask = 6 # fmt:off cls.input_tensor = torch.Tensor( [[[ 0.7049, 0.0305, -0.8542, 0.5388], [-0.5265, -1.3320, 1.5451, 0.4086], [-2.0546, 0.5259, 0.5995, -0.4078], [ 0.4530, -0.3918, 2.1403, -0.2062], [-0.0984, 0.4855, 0.7076, 0.0431], [-0.4394, -0.6761, 1.7389, -0.9423]]] ) # fmt:on def test_forward_causal_self_attn_and_no_xattn(self): set_seed(0) model = Transformer( n_layers=self.n_layers, d_model=self.d_model, d_ffn=self.d_ffn, sa_n_heads=self.sa_n_heads, kernel_size=self.kernel_size, p_dropout=self.p_dropout, p_dropout_out=self.p_dropout_out, has_xattn=False, is_causal=self.is_causal, max_length_causal_mask=self.max_length_causal_mask, ) # Check model init assert torch.isclose(torch.mean(model.layers[0].pos_ff.proj.conv.weight), torch.tensor(0.0), atol=1e-2) assert torch.isclose(torch.std(model.layers[0].pos_ff.proj.conv.weight), torch.tensor(0.02), atol=1e-2) assert torch.isclose(torch.mean(model.layers[0].pos_ff.o_net.conv.weight), torch.tensor(0.0), atol=1e-2) assert torch.isclose( torch.std(model.layers[0].pos_ff.o_net.conv.weight), torch.tensor(0.02 / math.sqrt(2.0)), atol=1e-3 ) mask_tensor = torch.ones(1, self.max_length_causal_mask).bool() with torch.no_grad(): output_dict = model(x=self.input_tensor, x_mask=mask_tensor) # fmt:off expected_output_tensor = { 'output': torch.Tensor( [[[0.7047, 0.0305, -0.8555, 0.5402], [-0.5192, -1.3324, 1.5455, 0.4148], [-2.0593, 0.5290, 0.5969, -0.4101], [0.4517, -0.3968, 2.1392, -0.2041], [-0.1019, 0.4854, 0.7077, 0.0431], [-0.4458, -0.6789, 1.7447, -0.9447]]] ), 'attn_probabilities': [ { 'self_attn_probabilities': [ torch.Tensor( [[[[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000], [0.4998, 0.5002, 0.0000, 0.0000, 0.0000, 0.0000], [0.3337, 0.3333, 0.3330, 0.0000, 0.0000, 0.0000], [0.2498, 0.2500, 0.2501, 0.2501, 0.0000, 0.0000], [0.2002, 0.2001, 0.2000, 0.1999, 0.1999, 0.0000], [0.1666, 0.1666, 0.1667, 0.1667, 0.1667, 0.1667]], [[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000], [0.5005, 0.4995, 0.0000, 0.0000, 0.0000, 0.0000], [0.3332, 0.3331, 0.3336, 0.0000, 0.0000, 0.0000], [0.2507, 0.2494, 0.2510, 0.2489, 0.0000, 0.0000], [0.2002, 0.1995, 0.2006, 0.1994, 0.2003, 0.0000], [0.1671, 0.1663, 0.1674, 0.1660, 0.1670, 0.1662]]]] ), torch.Tensor( [[[[-3.4823e-04, float('-inf'), float('-inf'), float('-inf'), float('-inf'), float('-inf')], [-7.0210e-04, -5.6984e-05, float('-inf'), float('-inf'), float('-inf'), float('-inf')], [1.2551e-03, 1.1431e-04, -5.9334e-04, float('-inf'), float('-inf'), float('-inf')], [-8.1514e-04, 2.6650e-05, 5.2952e-04, 5.6903e-04, float('-inf'), float('-inf')], [8.0150e-04, 1.4366e-04, -2.7793e-04, -7.0636e-04, -8.2140e-04, float('-inf')], [-4.6137e-04, 6.2648e-05, 3.6768e-04, 2.8095e-04, 4.9188e-04, 3.6888e-04]], [[-7.5861e-04, float('-inf'), float('-inf'), float('-inf'), float('-inf'), float('-inf')], [8.0373e-04, -1.2745e-03, float('-inf'), float('-inf'), float('-inf'), float('-inf')], [-4.5038e-04, -6.7648e-04, 7.5978e-04, float('-inf'), float('-inf'), float('-inf')], [1.5376e-03, -3.6984e-03, 3.0423e-03, -5.4870e-03, float('-inf'), float('-inf')], [3.7014e-04, -2.7010e-03, 2.4310e-03, -3.2604e-03, 9.3840e-04, float('-inf')], [1.3868e-03, -3.8372e-03, 3.2144e-03, -5.4860e-03, 5.0273e-04, -4.4343e-03]]]] ), ], 'cross_attn_probabilities': None, } ], } # fmt:on assert output_dict["output"].shape == expected_output_tensor["output"].shape assert torch.allclose(output_dict["output"], expected_output_tensor["output"], atol=1e-4) for i in range(2): assert torch.allclose( output_dict["attn_probabilities"][0]["self_attn_probabilities"][i], expected_output_tensor["attn_probabilities"][0]["self_attn_probabilities"][i], atol=1e-4, ) assert output_dict["attn_probabilities"][0]["cross_attn_probabilities"] is None def test_forward_causal_self_attn_and_has_xattn(self): set_seed(0) model = Transformer( n_layers=2, d_model=self.d_model, d_ffn=self.d_ffn, sa_n_heads=self.sa_n_heads, kernel_size=self.kernel_size, p_dropout=self.p_dropout, p_dropout_out=self.p_dropout_out, has_xattn=True, xa_d_memory=4, xa_n_heads=2, is_causal=self.is_causal, max_length_causal_mask=self.max_length_causal_mask, ) # fmt:off cond = [ # shape (1, 3, 4) torch.Tensor( [[[-0.7475, 1.1461, 0.7300, 1.4471], [ 1.8744, -0.1654, 1.2418, -1.6983], [-0.3123, 0.2320, 0.7457, 1.9868]]] ), # shape (1, 5, 4) torch.Tensor( [[[-0.6683, -1.2178, 1.3696, 0.9941], [ 0.0297, -0.1616, 0.1891, 0.0580], [-1.0771, 0.2547, -1.4023, 0.0971], [ 1.1132, 0.6311, -0.1449, 0.2351], [ 0.8920, 2.3663, 0.2248, -0.7298]]] ) ] # fmt:on cond_mask = [torch.ones(1, cond[0].shape[1]).bool(), torch.ones(1, cond[1].shape[1]).bool()] mask_tensor = torch.ones(1, self.max_length_causal_mask).bool() multi_encoder_mapping = [0, 1] with torch.no_grad(): output_dict = model( x=self.input_tensor, x_mask=mask_tensor, cond=cond, cond_mask=cond_mask, multi_encoder_mapping=multi_encoder_mapping, ) # fmt:off expected_output = { 'output': torch.Tensor( [[[0.7043, 0.0288, -0.8547, 0.5384], [-0.5283, -1.3311, 1.5429, 0.4083], [-2.0560, 0.5259, 0.6020, -0.4099], [0.4554, -0.3829, 2.1433, -0.2036], [-0.0986, 0.4794, 0.7067, 0.0432], [-0.4392, -0.6772, 1.7428, -0.9393]]] ), 'attn_probabilities': [ { 'self_attn_probabilities': [ torch.Tensor( [[[[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000], [0.4989, 0.5011, 0.0000, 0.0000, 0.0000, 0.0000], [0.3331, 0.3332, 0.3336, 0.0000, 0.0000, 0.0000], [0.2495, 0.2496, 0.2504, 0.2505, 0.0000, 0.0000], [0.1998, 0.1994, 0.2002, 0.2001, 0.2005, 0.0000], [0.1662, 0.1662, 0.1668, 0.1668, 0.1671, 0.1669]], [[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000], [0.5008, 0.4992, 0.0000, 0.0000, 0.0000, 0.0000], [0.3334, 0.3336, 0.3331, 0.0000, 0.0000, 0.0000], [0.2504, 0.2500, 0.2496, 0.2500, 0.0000, 0.0000], [0.2001, 0.2002, 0.2000, 0.1999, 0.1998, 0.0000], [0.1670, 0.1667, 0.1665, 0.1667, 0.1665, 0.1666]]]] ), ], 'cross_attn_probabilities': [ torch.Tensor( [[[[0.3331, 0.3336, 0.3334], [0.3335, 0.3331, 0.3334], [0.3336, 0.3331, 0.3332], [0.3335, 0.3332, 0.3334], [0.3336, 0.3331, 0.3333], [0.3335, 0.3332, 0.3333]], [[0.3333, 0.3335, 0.3332], [0.3334, 0.3335, 0.3331], [0.3333, 0.3330, 0.3337], [0.3334, 0.3335, 0.3332], [0.3333, 0.3331, 0.3336], [0.3334, 0.3334, 0.3333]]]] ) ] }, { 'self_attn_probabilities': [ torch.Tensor( [[[[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000], [0.5005, 0.4995, 0.0000, 0.0000, 0.0000, 0.0000], [0.3336, 0.3330, 0.3334, 0.0000, 0.0000, 0.0000], [0.2503, 0.2499, 0.2498, 0.2500, 0.0000, 0.0000], [0.2002, 0.1999, 0.2000, 0.2000, 0.2000, 0.0000], [0.1669, 0.1666, 0.1666, 0.1667, 0.1666, 0.1666]], [[1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000], [0.5001, 0.4999, 0.0000, 0.0000, 0.0000, 0.0000], [0.3329, 0.3330, 0.3340, 0.0000, 0.0000, 0.0000], [0.2499, 0.2498, 0.2505, 0.2499, 0.0000, 0.0000], [0.1997, 0.1997, 0.2003, 0.2000, 0.2004, 0.0000], [0.1665, 0.1664, 0.1669, 0.1666, 0.1669, 0.1667]]]] ), ], 'cross_attn_probabilities': [ torch.Tensor( [[[[0.1999, 0.1998, 0.2002, 0.1999, 0.2001], [0.2000, 0.1997, 0.2004, 0.1998, 0.2002], [0.2001, 0.2000, 0.2001, 0.1999, 0.2000], [0.2000, 0.2001, 0.1998, 0.2001, 0.1999], [0.2001, 0.2002, 0.1998, 0.2001, 0.1998], [0.2000, 0.2002, 0.1998, 0.2001, 0.1999]], [[0.1998, 0.1998, 0.2001, 0.2004, 0.2000], [0.2003, 0.2003, 0.1998, 0.1995, 0.2001], [0.2003, 0.2003, 0.1998, 0.1995, 0.2001], [0.2001, 0.2001, 0.2000, 0.1998, 0.2000], [0.2002, 0.2001, 0.1999, 0.1997, 0.2000], [0.2002, 0.2001, 0.2000, 0.1997, 0.2000]]]] ), ], } ], } # fmt:on assert torch.allclose(output_dict["output"], expected_output["output"], atol=1e-4) for i in range(2): assert torch.allclose( output_dict["attn_probabilities"][i]["self_attn_probabilities"][0], expected_output["attn_probabilities"][i]["self_attn_probabilities"][0], atol=1e-4, ) assert torch.allclose( output_dict["attn_probabilities"][i]["cross_attn_probabilities"][0], expected_output["attn_probabilities"][i]["cross_attn_probabilities"][0], atol=1e-4, )