71 lines
1.6 KiB
Python
71 lines
1.6 KiB
Python
import torch
|
|
import pytest
|
|
|
|
from pina.model.block import LowRankBlock
|
|
from pina import LabelTensor
|
|
|
|
|
|
input_dimensions = 2
|
|
embedding_dimenion = 1
|
|
rank = 4
|
|
inner_size = 20
|
|
n_layers = 2
|
|
func = torch.nn.Tanh
|
|
bias = True
|
|
|
|
|
|
def test_constructor():
|
|
LowRankBlock(
|
|
input_dimensions=input_dimensions,
|
|
embedding_dimenion=embedding_dimenion,
|
|
rank=rank,
|
|
inner_size=inner_size,
|
|
n_layers=n_layers,
|
|
func=func,
|
|
bias=bias,
|
|
)
|
|
|
|
|
|
def test_constructor_wrong():
|
|
with pytest.raises(ValueError):
|
|
LowRankBlock(
|
|
input_dimensions=input_dimensions,
|
|
embedding_dimenion=embedding_dimenion,
|
|
rank=0.5,
|
|
inner_size=inner_size,
|
|
n_layers=n_layers,
|
|
func=func,
|
|
bias=bias,
|
|
)
|
|
|
|
|
|
def test_forward():
|
|
block = LowRankBlock(
|
|
input_dimensions=input_dimensions,
|
|
embedding_dimenion=embedding_dimenion,
|
|
rank=rank,
|
|
inner_size=inner_size,
|
|
n_layers=n_layers,
|
|
func=func,
|
|
bias=bias,
|
|
)
|
|
data = LabelTensor(torch.rand(10, 30, 3), labels=["x", "y", "u"])
|
|
block(data.extract("u"), data.extract(["x", "y"]))
|
|
|
|
|
|
def test_backward():
|
|
block = LowRankBlock(
|
|
input_dimensions=input_dimensions,
|
|
embedding_dimenion=embedding_dimenion,
|
|
rank=rank,
|
|
inner_size=inner_size,
|
|
n_layers=n_layers,
|
|
func=func,
|
|
bias=bias,
|
|
)
|
|
data = LabelTensor(torch.rand(10, 30, 3), labels=["x", "y", "u"])
|
|
data.requires_grad_(True)
|
|
out = block(data.extract("u"), data.extract(["x", "y"]))
|
|
loss = out.mean()
|
|
loss.backward()
|