54 lines
1.4 KiB
Python
54 lines
1.4 KiB
Python
import torch
|
|
import pytest
|
|
from pina.model.block import PirateNetBlock
|
|
|
|
data = torch.rand((20, 3))
|
|
|
|
|
|
@pytest.mark.parametrize("inner_size", [10, 20])
|
|
def test_constructor(inner_size):
|
|
|
|
PirateNetBlock(inner_size=inner_size, activation=torch.nn.Tanh)
|
|
|
|
# Should fail if inner_size is negative
|
|
with pytest.raises(AssertionError):
|
|
PirateNetBlock(inner_size=-1, activation=torch.nn.Tanh)
|
|
|
|
|
|
@pytest.mark.parametrize("inner_size", [10, 20])
|
|
def test_forward(inner_size):
|
|
|
|
model = PirateNetBlock(inner_size=inner_size, activation=torch.nn.Tanh)
|
|
|
|
# Create dummy embedding
|
|
dummy_embedding = torch.nn.Linear(data.shape[1], inner_size)
|
|
x = dummy_embedding(data)
|
|
|
|
# Create dummy U and V tensors
|
|
U = torch.rand((data.shape[0], inner_size))
|
|
V = torch.rand((data.shape[0], inner_size))
|
|
|
|
output_ = model(x, U, V)
|
|
assert output_.shape == (data.shape[0], inner_size)
|
|
|
|
|
|
@pytest.mark.parametrize("inner_size", [10, 20])
|
|
def test_backward(inner_size):
|
|
|
|
model = PirateNetBlock(inner_size=inner_size, activation=torch.nn.Tanh)
|
|
data.requires_grad_()
|
|
|
|
# Create dummy embedding
|
|
dummy_embedding = torch.nn.Linear(data.shape[1], inner_size)
|
|
x = dummy_embedding(data)
|
|
|
|
# Create dummy U and V tensors
|
|
U = torch.rand((data.shape[0], inner_size))
|
|
V = torch.rand((data.shape[0], inner_size))
|
|
|
|
output_ = model(x, U, V)
|
|
|
|
loss = torch.mean(output_)
|
|
loss.backward()
|
|
assert data.grad.shape == data.shape
|