def jac_vec_prod(self, vec) -> Tensor:
"""Product of input-output-Jacobian and a vector.
Args:
vec: vector
Returns:
product
"""
input, output, _ = self.problem.forward_pass(input_requires_grad=True)
return jacobian_vector_product(output, input, vec)[0]
def test_input_jacobian(self):
"""Test multiplication by the input Jacobian.
Compare with result of R-operator.
"""
# create input
cvp_in = self._create_input()
cvp_in.requires_grad = True
cvp_layer = self._create_cvp_layer()
# skip for Sequential:
if isinstance(cvp_layer, (Flatten, CVPSequential)):
return
cvp_out = cvp_layer(cvp_in)
for _ in range(self.NUM_HVP):
v = torch.randn(cvp_in.numel(),
requires_grad=False).to(self.DEVICE)
Jv = cvp_layer._input_jacobian(v)
# compute via R-operator
(result, ) = jacobian_vector_product(cvp_out, cvp_in,
v.view(cvp_in.size()))
assert torch.allclose(Jv,
result.view(-1),
atol=self.ATOL,
rtol=self.RTOL)
def f():
r = jacobian_vector_product(out, X, vout, detach=False)[0].contiguous()
if vout.is_cuda:
torch.cuda.synchronize()
return r
def param_jac_vec_prod(self, name, vec):
input, output, named_params = self.problem.forward_pass()
param = named_params[name]
return jacobian_vector_product(output, param, vec)[0]
def jac_vec_prod(self, vec):
input, output, _ = self.problem.forward_pass(input_requires_grad=True)
return jacobian_vector_product(output, input, vec)[0]