(self)
| 2029 | self.assertEqual(out.shape, x.shape) |
| 2030 | |
| 2031 | def test_transformer_decoder_layer(self): |
| 2032 | dims = 32 |
| 2033 | num_heads = 4 |
| 2034 | x = mx.random.normal(shape=(2, 5, dims)) |
| 2035 | memory = mx.random.normal(shape=(2, 8, dims)) |
| 2036 | |
| 2037 | # Test norm_first=True (default) |
| 2038 | layer = nn.TransformerDecoderLayer(dims=dims, num_heads=num_heads) |
| 2039 | out = layer(x, memory, x_mask=None, memory_mask=None) |
| 2040 | self.assertEqual(out.shape, x.shape) |
| 2041 | |
| 2042 | # Test norm_first=False |
| 2043 | layer = nn.TransformerDecoderLayer( |
| 2044 | dims=dims, num_heads=num_heads, norm_first=False |
| 2045 | ) |
| 2046 | out = layer(x, memory, x_mask=None, memory_mask=None) |
| 2047 | self.assertEqual(out.shape, x.shape) |
| 2048 | |
| 2049 | # Test with masks |
| 2050 | x_mask = nn.MultiHeadAttention.create_additive_causal_mask(5) |
| 2051 | out = layer(x, memory, x_mask=x_mask, memory_mask=None) |
| 2052 | self.assertEqual(out.shape, x.shape) |
| 2053 | |
| 2054 | # Test with custom mlp_dims |
| 2055 | layer = nn.TransformerDecoderLayer(dims=dims, num_heads=num_heads, mlp_dims=64) |
| 2056 | out = layer(x, memory, x_mask=None, memory_mask=None) |
| 2057 | self.assertEqual(out.shape, x.shape) |
| 2058 | |
| 2059 | def test_transformer_encoder(self): |
| 2060 | encoder = nn.TransformerEncoder(num_layers=2, dims=32, num_heads=4) |
nothing calls this directly
no test coverage detected