def test_diag(self):
# test that it builds a matrix with given diagonal when using
# vector inputs
x = vector()
y = diag(x)
assert y.owner.op.__class__ == AllocDiag
# test that it extracts the diagonal when using matrix input
x = matrix()
y = extract_diag(x)
assert y.owner.op.__class__ == ExtractDiag
def grad(self, inp, cost_grad):
"""
Notes
-----
The gradient is currently implemented for matrices only.
"""
a, val = inp
grad = cost_grad[0]
if a.dtype.startswith("complex"):
return [None, None]
elif a.ndim > 2:
raise NotImplementedError("%s: gradient is currently implemented"
" for matrices only" %
self.__class__.__name__)
wr_a = fill_diagonal(grad, 0) # valid for any number of dimensions
# diag is only valid for matrices
wr_val = nlinalg.diag(grad).sum()
return [wr_a, wr_val]