def test_mul_adding_constant_mul(self):
mat1 = make_random_mat(20, rank=4, batch_size=5)
mat2 = make_random_mat(20, rank=4, batch_size=5)
mat3 = make_random_mat(20, rank=4, batch_size=5)
const = torch.ones(1, requires_grad=True)
mat1_copy = mat1.clone().detach().requires_grad_(True)
mat2_copy = mat2.clone().detach().requires_grad_(True)
mat3_copy = mat3.clone().detach().requires_grad_(True)
const_copy = const.clone().detach().requires_grad_(True)
# Forward
res = MulLazyTensor(RootLazyTensor(mat1), RootLazyTensor(mat2), RootLazyTensor(mat3))
res = res * const
actual = (
prod(
[
mat1_copy.matmul(mat1_copy.transpose(-1, -2)),
mat2_copy.matmul(mat2_copy.transpose(-1, -2)),
mat3_copy.matmul(mat3_copy.transpose(-1, -2)),
]
)
* const_copy
)
self.assertLess(torch.max(((res.evaluate() - actual) / actual).abs()), 0.01)
# Forward
res = MulLazyTensor(RootLazyTensor(mat1), RootLazyTensor(mat2), RootLazyTensor(mat3))
res = res * 2.5
actual = (
prod(
[
mat1_copy.matmul(mat1_copy.transpose(-1, -2)),
mat2_copy.matmul(mat2_copy.transpose(-1, -2)),
mat3_copy.matmul(mat3_copy.transpose(-1, -2)),
]
)
* 2.5
)
self.assertLess(torch.max(((res.evaluate() - actual) / actual).abs()), 0.01)
def test_mul_adding_another_variable(self):
mat1 = make_random_mat(20, rank=4, batch_size=5)
mat2 = make_random_mat(20, rank=4, batch_size=5)
mat3 = make_random_mat(20, rank=4, batch_size=5)
mat1_copy = mat1.clone().detach().requires_grad_(True)
mat2_copy = mat2.clone().detach().requires_grad_(True)
mat3_copy = mat3.clone().detach().requires_grad_(True)
# Forward
res = MulLazyTensor(RootLazyTensor(mat1), RootLazyTensor(mat2))
res = res * RootLazyTensor(mat3)
actual = prod(
[
mat1_copy.matmul(mat1_copy.transpose(-1, -2)),
mat2_copy.matmul(mat2_copy.transpose(-1, -2)),
mat3_copy.matmul(mat3_copy.transpose(-1, -2)),
]
)
self.assertLess(torch.max(((res.evaluate() - actual) / actual).abs()), 0.01)