def sum(a: Kronecker, axis=None):
if axis is None:
return B.sum(a.left) * B.sum(a.right)
elif axis == 0:
return B.kron(B.sum(a.left, axis=0), B.sum(a.right, axis=0))
elif axis == 1:
return B.kron(B.sum(a.left, axis=1), B.sum(a.right, axis=1))
else:
_raise(axis)
def test_kron(check_lazy_shapes):
# Cannot test tensors of rank higher than two, because TensorFlows broadcasting
# behaviour does not allow that.
# Test full Kronecker product.
check_function(B.kron, (Tensor(2, 3), Tensor(4, 5)))
check_function(lambda x, y: B.kron(x, y, 0, 1), (Tensor(2, 3), Tensor(4, 5)))
check_function(lambda x, y: B.kron(x, y, -2, -1), (Tensor(2, 3), Tensor(4, 5)))
# Test Kronecker product over dimension 0.
check_function(lambda x, y: B.kron(x, y, 0), (Tensor(2, 3), Tensor(4, 3)))
check_function(lambda x, y: B.kron(x, y, -2), (Tensor(2, 3), Tensor(4, 3)))
# Test Kronecker product over dimension 1.
check_function(lambda x, y: B.kron(x, y, 1), (Tensor(2, 3), Tensor(2, 5)))
check_function(lambda x, y: B.kron(x, y, -1), (Tensor(2, 3), Tensor(2, 5)))
# Test checking of shapes.
with pytest.raises(ValueError):
B.kron(Tensor(2).np(), Tensor(4, 5).np())
with pytest.raises(ValueError):
B.kron(Tensor(2, 3).np(), Tensor(4, 5).np(), 1)
def dense(a: Kronecker):
if a.dense is None:
a.dense = B.kron(B.dense(a.left), B.dense(a.right))
return a.dense
def diag(a: Kronecker):
return B.kron(B.diag(a.left), B.diag(a.right))
def test_dense_kron(kron1):
approx(B.dense(kron1), B.kron(B.dense(kron1.left), B.dense(kron1.right)))
_check_cache(kron1)