(self)
| 2063 | self.assertEqual(out.shape, x.shape) |
| 2064 | |
| 2065 | def test_transformer_decoder(self): |
| 2066 | decoder = nn.TransformerDecoder(num_layers=2, dims=32, num_heads=4) |
| 2067 | x = mx.random.normal(shape=(2, 5, 32)) |
| 2068 | memory = mx.random.normal(shape=(2, 8, 32)) |
| 2069 | out = decoder(x, memory, x_mask=None, memory_mask=None) |
| 2070 | self.assertEqual(out.shape, x.shape) |
| 2071 | |
| 2072 | def test_transformer(self): |
| 2073 | model = nn.Transformer( |
nothing calls this directly
no outgoing calls
no test coverage detected