38 lines
1.0 KiB
Python
38 lines
1.0 KiB
Python
"""Tests for ScrollCVAE model."""
|
|
from __future__ import annotations
|
|
|
|
import torch
|
|
import pytest
|
|
|
|
from tools.scroll.models import ScrollCVAE
|
|
|
|
|
|
class TestScrollCVAEForward:
|
|
@pytest.fixture
|
|
def model(self):
|
|
return ScrollCVAE(seq_len=32, latent_dim=16, hidden=64, cond_dim=7)
|
|
|
|
def test_output_shapes(self, model):
|
|
batch = 4
|
|
seq = torch.randn(batch, 32, 2)
|
|
cond = torch.randn(batch, 7)
|
|
recon, mu, logvar = model(seq, cond)
|
|
assert recon.shape == (batch, 32, 2)
|
|
assert mu.shape == (batch, 16)
|
|
assert logvar.shape == (batch, 16)
|
|
|
|
def test_decode_shape(self, model):
|
|
z = torch.randn(4, 16)
|
|
cond = torch.randn(4, 7)
|
|
out = model.decode(z, cond)
|
|
assert out.shape == (4, 32, 2)
|
|
|
|
def test_decode_deterministic(self, model):
|
|
model.eval()
|
|
z = torch.randn(1, 16)
|
|
cond = torch.randn(1, 7)
|
|
with torch.no_grad():
|
|
out1 = model.decode(z, cond)
|
|
out2 = model.decode(z, cond)
|
|
torch.testing.assert_close(out1, out2)
|