@@ -0,0 +1,1554 @@
+import sys%0Aimport pytest%0Afrom numpy.testing import assert_equal%0Aimport torch%0A%0Asys.path.append(%22../../../%22)%0A%0Afrom pycroscopy.learn import models%0A%0A%0A@pytest.mark.parametrize(%22dim, size%22, %5B(1, %5B8%5D), (2, %5B8, 8%5D), (3, %5B8, 8, 8%5D)%5D)%0Adef test_autoencoder_output(dim, size):%0A    input_dim = (1, *size)%0A    x = torch.randn(2, *input_dim)%0A    ae = models.AutoEncoder(input_dim, 2, %5B1, 1%5D)%0A    out = ae(x)%0A    assert_equal(input_dim, out.shape%5B1:%5D)%0A%0A%0A@pytest.mark.parametrize(%22zdim%22, %5B1, 2, 5%5D)%0A@pytest.mark.parametrize(%22dim, size%22, %5B(1, %5B8%5D), (2, %5B8, 8%5D), (3, %5B8, 8, 8%5D)%5D)%0Adef test_autoencoder_encoding(zdim, dim, size):%0A    input_dim = (1, *size)%0A    x = torch.randn(2, *input_dim)%0A    ae = models.AutoEncoder(input_dim, zdim, %5B1, 1%5D)%0A    z = ae.encode(x)%0A    assert_equal(zdim, z.shape%5B-1%5D)%0A%0A%0A@pytest.mark.parametrize(%22zdim%22, %5B1, 2, 5%5D)%0A@pytest.mark.parametrize(%22dim, size%22, %5B(1, %5B8%5D), (2, %5B8, 8%5D), (3, %5B8, 8, 8%5D)%5D)%0Adef test_autoencoder_decoding(zdim, dim, size):%0A    input_dim = (1, *size)%0A    x = torch.randn(2, *input_dim)%0A    ae = models.AutoEncoder(input_dim, zdim, %5B1, 1%5D)%0A    z = torch.randn(zdim)%0A    x_ = ae.decode(z)%0A    assert_equal(x_.shape%5B1:%5D, x.shape%5B1:%5D)%0A%0A%0A@pytest.mark.parametrize(%22dim, size%22, %5B(1, %5B8%5D), (2, %5B8, 8%5D), (3, %5B8, 8, 8%5D)%5D)%0Adef test_autoencoder_decode_grid(dim, size):%0A    input_dim = (1, *size)%0A    grid_spacing = 4%0A    x = torch.randn(2, *input_dim)%0A    ae = models.AutoEncoder(input_dim, 2, %5B1, 1%5D)%0A    grid = ae.decode_grid(grid_spacing)%0A    assert_equal(grid.shape%5B0%5D, grid_spacing**2)%0A    assert_equal(grid.shape%5B1:%5D, x.shape%5B2:%5D)

Add tests for autoencoder models