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]