| from typing import List, Tuple |
| |
| import torch |
| from torch._vmap_internals import _vmap |
| from . import forward_ad as fwAD |
| |
| __all__ = ["vjp", "jvp", "jacobian", "hessian", "hvp", "vhp"] |
| |
| # Utility functions |
| |
| |
| def _as_tuple_nocheck(x): |
| if isinstance(x, tuple): |
| return x |
| elif isinstance(x, list): |
| return tuple(x) |
| else: |
| return (x,) |
| |
| |
| def _as_tuple(inp, arg_name=None, fn_name=None): |
| # Ensures that inp is a tuple of Tensors |
| # Returns whether or not the original inp was a tuple and the tupled version of the input |
| if arg_name is None and fn_name is None: |
| return _as_tuple_nocheck(inp) |
| |
| is_inp_tuple = True |
| if not isinstance(inp, tuple): |
| inp = (inp,) |
| is_inp_tuple = False |
| |
| for i, el in enumerate(inp): |
| if not isinstance(el, torch.Tensor): |
| if is_inp_tuple: |
| raise TypeError( |
| f"The {arg_name} given to {fn_name} must be either a Tensor or a tuple of Tensors but the" |
| f" value at index {i} has type {type(el)}." |
| ) |
| else: |
| raise TypeError( |
| f"The {arg_name} given to {fn_name} must be either a Tensor or a tuple of Tensors but the" |
| f" given {arg_name} has type {type(el)}." |
| ) |
| |
| return is_inp_tuple, inp |
| |
| |
| def _tuple_postprocess(res, to_unpack): |
| # Unpacks a potentially nested tuple of Tensors |
| # to_unpack should be a single boolean or a tuple of two booleans. |
| # It is used to: |
| # - invert _as_tuple when res should match the inp given to _as_tuple |
| # - optionally remove nesting of two tuples created by multiple calls to _as_tuple |
| if isinstance(to_unpack, tuple): |
| assert len(to_unpack) == 2 |
| if not to_unpack[1]: |
| res = tuple(el[0] for el in res) |
| if not to_unpack[0]: |
| res = res[0] |
| else: |
| if not to_unpack: |
| res = res[0] |
| return res |
| |
| |
| def _grad_preprocess(inputs, create_graph, need_graph): |
| # Preprocess the inputs to make sure they require gradient |
| # inputs is a tuple of Tensors to preprocess |
| # create_graph specifies if the user wants gradients to flow back to the Tensors in inputs |
| # need_graph specifies if we internally want gradients to flow back to the Tensors in res |
| # Note that we *always* create a new Tensor object to be able to see the difference between |
| # inputs given as arguments and the same Tensors automatically captured by the user function. |
| # Check this issue for more details on how that can happen: https://github.com/pytorch/pytorch/issues/32576 |
| res = [] |
| for inp in inputs: |
| if create_graph and inp.requires_grad: |
| # Create at least a new Tensor object in a differentiable way |
| if not inp.is_sparse: |
| # Use .view_as() to get a shallow copy |
| res.append(inp.view_as(inp)) |
| else: |
| # We cannot use view for sparse Tensors so we clone |
| res.append(inp.clone()) |
| else: |
| res.append(inp.detach().requires_grad_(need_graph)) |
| return tuple(res) |
| |
| |
| def _grad_postprocess(inputs, create_graph): |
| # Postprocess the generated Tensors to avoid returning Tensors with history when the user did not |
| # request it. |
| if isinstance(inputs[0], torch.Tensor): |
| if not create_graph: |
| return tuple(inp.detach() for inp in inputs) |
| else: |
| return inputs |
| else: |
| return tuple(_grad_postprocess(inp, create_graph) for inp in inputs) |
| |
| |
| def _validate_v(v, other, is_other_tuple): |
| # This assumes that other is the correct shape, and v should match |
| # Both are assumed to be tuples of Tensors |
| if len(other) != len(v): |
| if is_other_tuple: |
| raise RuntimeError( |
| f"v is a tuple of invalid length: should be {len(other)} but got {len(v)}." |
| ) |
| else: |
| raise RuntimeError("The given v should contain a single Tensor.") |
| |
| for idx, (el_v, el_other) in enumerate(zip(v, other)): |
| if el_v.size() != el_other.size(): |
| prepend = "" |
| if is_other_tuple: |
| prepend = f"Entry {idx} in " |
| raise RuntimeError( |
| f"{prepend}v has invalid size: should be {el_other.size()} but got {el_v.size()}." |
| ) |
| |
| |
| def _check_requires_grad(inputs, input_type, strict): |
| # Used to make all the necessary checks to raise nice errors in strict mode. |
| if not strict: |
| return |
| |
| if input_type not in ["outputs", "grad_inputs", "jacobian", "hessian"]: |
| raise RuntimeError("Invalid input_type to _check_requires_grad") |
| for i, inp in enumerate(inputs): |
| if inp is None: |
| # This can only be reached for grad_inputs. |
| raise RuntimeError( |
| f"The output of the user-provided function is independent of input {i}." |
| " This is not allowed in strict mode." |
| ) |
| if not inp.requires_grad: |
| if input_type == "hessian": |
| raise RuntimeError( |
| f"The hessian of the user-provided function with respect to input {i}" |
| " is independent of the input. This is not allowed in strict mode." |
| " You should ensure that your function is thrice differentiable and that" |
| " the hessian depends on the inputs." |
| ) |
| elif input_type == "jacobian": |
| raise RuntimeError( |
| "While computing the hessian, found that the jacobian of the user-provided" |
| f" function with respect to input {i} is independent of the input. This is not" |
| " allowed in strict mode. You should ensure that your function is twice" |
| " differentiable and that the jacobian depends on the inputs (this would be" |
| " violated by a linear function for example)." |
| ) |
| elif input_type == "grad_inputs": |
| raise RuntimeError( |
| f"The gradient with respect to input {i} is independent of the inputs of the" |
| " user-provided function. This is not allowed in strict mode." |
| ) |
| else: |
| raise RuntimeError( |
| f"Output {i} of the user-provided function does not require gradients." |
| " The outputs must be computed in a differentiable manner from the input" |
| " when running in strict mode." |
| ) |
| |
| |
| def _autograd_grad( |
| outputs, |
| inputs, |
| grad_outputs=None, |
| create_graph=False, |
| retain_graph=None, |
| is_grads_batched=False, |
| ): |
| # Version of autograd.grad that accepts `None` in outputs and do not compute gradients for them. |
| # This has the extra constraint that inputs has to be a tuple |
| assert isinstance(outputs, tuple) |
| if grad_outputs is None: |
| grad_outputs = (None,) * len(outputs) |
| assert isinstance(grad_outputs, tuple) |
| assert len(outputs) == len(grad_outputs) |
| |
| new_outputs: Tuple[torch.Tensor, ...] = tuple() |
| new_grad_outputs: Tuple[torch.Tensor, ...] = tuple() |
| for out, grad_out in zip(outputs, grad_outputs): |
| if out is not None and out.requires_grad: |
| new_outputs += (out,) |
| new_grad_outputs += (grad_out,) |
| |
| if len(new_outputs) == 0: |
| # No differentiable output, we don't need to call the autograd engine |
| return (None,) * len(inputs) |
| else: |
| return torch.autograd.grad( |
| new_outputs, |
| inputs, |
| new_grad_outputs, |
| allow_unused=True, |
| create_graph=create_graph, |
| retain_graph=retain_graph, |
| is_grads_batched=is_grads_batched, |
| ) |
| |
| |
| def _fill_in_zeros(grads, refs, strict, create_graph, stage): |
| # Used to detect None in the grads and depending on the flags, either replace them |
| # with Tensors full of 0s of the appropriate size based on the refs or raise an error. |
| # strict and create graph allow us to detect when it is appropriate to raise an error |
| # stage gives us information of which backward call we consider to give good error message |
| if stage not in ["back", "back_trick", "double_back", "double_back_trick"]: |
| raise RuntimeError(f"Invalid stage argument '{stage}' to _fill_in_zeros") |
| |
| res: Tuple[torch.Tensor, ...] = tuple() |
| for i, grads_i in enumerate(grads): |
| if grads_i is None: |
| if strict: |
| if stage == "back": |
| raise RuntimeError( |
| "The output of the user-provided function is independent of " |
| f"input {i}. This is not allowed in strict mode." |
| ) |
| elif stage == "back_trick": |
| raise RuntimeError( |
| f"The gradient with respect to the input is independent of entry {i}" |
| " in the grad_outputs when using the double backward trick to compute" |
| " forward mode gradients. This is not allowed in strict mode." |
| ) |
| elif stage == "double_back": |
| raise RuntimeError( |
| "The jacobian of the user-provided function is independent of " |
| f"input {i}. This is not allowed in strict mode." |
| ) |
| else: |
| raise RuntimeError( |
| "The hessian of the user-provided function is independent of " |
| f"entry {i} in the grad_jacobian. This is not allowed in strict " |
| "mode as it prevents from using the double backward trick to " |
| "replace forward mode AD." |
| ) |
| |
| grads_i = torch.zeros_like(refs[i]) |
| else: |
| if strict and create_graph and not grads_i.requires_grad: |
| if "double" not in stage: |
| raise RuntimeError( |
| "The jacobian of the user-provided function is independent of " |
| f"input {i}. This is not allowed in strict mode when create_graph=True." |
| ) |
| else: |
| raise RuntimeError( |
| "The hessian of the user-provided function is independent of " |
| f"input {i}. This is not allowed in strict mode when create_graph=True." |
| ) |
| |
| res += (grads_i,) |
| |
| return res |
| |
| |
| # Public API |
| |
| |
| def vjp(func, inputs, v=None, create_graph=False, strict=False): |
| r"""Function that computes the dot product between a vector ``v`` and the |
| Jacobian of the given function at the point given by the inputs. |
| |
| Args: |
| func (function): a Python function that takes Tensor inputs and returns |
| a tuple of Tensors or a Tensor. |
| inputs (tuple of Tensors or Tensor): inputs to the function ``func``. |
| v (tuple of Tensors or Tensor): The vector for which the vector |
| Jacobian product is computed. Must be the same size as the output |
| of ``func``. This argument is optional when the output of ``func`` |
| contains a single element and (if it is not provided) will be set |
| as a Tensor containing a single ``1``. |
| create_graph (bool, optional): If ``True``, both the output and result |
| will be computed in a differentiable way. Note that when ``strict`` |
| is ``False``, the result can not require gradients or be |
| disconnected from the inputs. Defaults to ``False``. |
| strict (bool, optional): If ``True``, an error will be raised when we |
| detect that there exists an input such that all the outputs are |
| independent of it. If ``False``, we return a Tensor of zeros as the |
| vjp for said inputs, which is the expected mathematical value. |
| Defaults to ``False``. |
| |
| Returns: |
| output (tuple): tuple with: |
| func_output (tuple of Tensors or Tensor): output of ``func(inputs)`` |
| |
| vjp (tuple of Tensors or Tensor): result of the dot product with |
| the same shape as the inputs. |
| |
| Example: |
| |
| >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_AUTOGRAD) |
| >>> def exp_reducer(x): |
| ... return x.exp().sum(dim=1) |
| >>> inputs = torch.rand(4, 4) |
| >>> v = torch.ones(4) |
| >>> # xdoctest: +IGNORE_WANT("non-deterministic") |
| >>> vjp(exp_reducer, inputs, v) |
| (tensor([5.7817, 7.2458, 5.7830, 6.7782]), |
| tensor([[1.4458, 1.3962, 1.3042, 1.6354], |
| [2.1288, 1.0652, 1.5483, 2.5035], |
| [2.2046, 1.1292, 1.1432, 1.3059], |
| [1.3225, 1.6652, 1.7753, 2.0152]])) |
| |
| >>> vjp(exp_reducer, inputs, v, create_graph=True) |
| (tensor([5.7817, 7.2458, 5.7830, 6.7782], grad_fn=<SumBackward1>), |
| tensor([[1.4458, 1.3962, 1.3042, 1.6354], |
| [2.1288, 1.0652, 1.5483, 2.5035], |
| [2.2046, 1.1292, 1.1432, 1.3059], |
| [1.3225, 1.6652, 1.7753, 2.0152]], grad_fn=<MulBackward0>)) |
| |
| >>> def adder(x, y): |
| ... return 2 * x + 3 * y |
| >>> inputs = (torch.rand(2), torch.rand(2)) |
| >>> v = torch.ones(2) |
| >>> vjp(adder, inputs, v) |
| (tensor([2.4225, 2.3340]), |
| (tensor([2., 2.]), tensor([3., 3.]))) |
| """ |
| |
| with torch.enable_grad(): |
| is_inputs_tuple, inputs = _as_tuple(inputs, "inputs", "vjp") |
| inputs = _grad_preprocess(inputs, create_graph=create_graph, need_graph=True) |
| |
| outputs = func(*inputs) |
| is_outputs_tuple, outputs = _as_tuple( |
| outputs, "outputs of the user-provided function", "vjp" |
| ) |
| _check_requires_grad(outputs, "outputs", strict=strict) |
| |
| if v is not None: |
| _, v = _as_tuple(v, "v", "vjp") |
| v = _grad_preprocess(v, create_graph=create_graph, need_graph=False) |
| _validate_v(v, outputs, is_outputs_tuple) |
| else: |
| if len(outputs) != 1 or outputs[0].nelement() != 1: |
| raise RuntimeError( |
| "The vector v can only be None if the " |
| "user-provided function returns " |
| "a single Tensor with a single element." |
| ) |
| |
| enable_grad = True if create_graph else torch.is_grad_enabled() |
| with torch.set_grad_enabled(enable_grad): |
| grad_res = _autograd_grad(outputs, inputs, v, create_graph=create_graph) |
| vjp = _fill_in_zeros(grad_res, inputs, strict, create_graph, "back") |
| |
| # Cleanup objects and return them to the user |
| outputs = _grad_postprocess(outputs, create_graph) |
| vjp = _grad_postprocess(vjp, create_graph) |
| |
| return _tuple_postprocess(outputs, is_outputs_tuple), _tuple_postprocess( |
| vjp, is_inputs_tuple |
| ) |
| |
| |
| def jvp(func, inputs, v=None, create_graph=False, strict=False): |
| r"""Function that computes the dot product between the Jacobian of |
| the given function at the point given by the inputs and a vector ``v``. |
| |
| Args: |
| func (function): a Python function that takes Tensor inputs and returns |
| a tuple of Tensors or a Tensor. |
| inputs (tuple of Tensors or Tensor): inputs to the function ``func``. |
| v (tuple of Tensors or Tensor): The vector for which the Jacobian |
| vector product is computed. Must be the same size as the input of |
| ``func``. This argument is optional when the input to ``func`` |
| contains a single element and (if it is not provided) will be set |
| as a Tensor containing a single ``1``. |
| create_graph (bool, optional): If ``True``, both the output and result |
| will be computed in a differentiable way. Note that when ``strict`` |
| is ``False``, the result can not require gradients or be |
| disconnected from the inputs. Defaults to ``False``. |
| strict (bool, optional): If ``True``, an error will be raised when we |
| detect that there exists an input such that all the outputs are |
| independent of it. If ``False``, we return a Tensor of zeros as the |
| jvp for said inputs, which is the expected mathematical value. |
| Defaults to ``False``. |
| |
| Returns: |
| output (tuple): tuple with: |
| func_output (tuple of Tensors or Tensor): output of ``func(inputs)`` |
| |
| jvp (tuple of Tensors or Tensor): result of the dot product with |
| the same shape as the output. |
| |
| Note: |
| ``autograd.functional.jvp`` computes the jvp by using the backward of |
| the backward (sometimes called the double backwards trick). This is not |
| the most performant way of computing the jvp. Please consider using |
| :func:`torch.func.jvp` or the |
| :ref:`low-level forward-mode AD API <forward-mode-ad>` instead. |
| |
| Example: |
| |
| >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_AUTOGRAD) |
| >>> def exp_reducer(x): |
| ... return x.exp().sum(dim=1) |
| >>> inputs = torch.rand(4, 4) |
| >>> v = torch.ones(4, 4) |
| >>> # xdoctest: +IGNORE_WANT("non-deterministic") |
| >>> jvp(exp_reducer, inputs, v) |
| (tensor([6.3090, 4.6742, 7.9114, 8.2106]), |
| tensor([6.3090, 4.6742, 7.9114, 8.2106])) |
| |
| >>> jvp(exp_reducer, inputs, v, create_graph=True) |
| (tensor([6.3090, 4.6742, 7.9114, 8.2106], grad_fn=<SumBackward1>), |
| tensor([6.3090, 4.6742, 7.9114, 8.2106], grad_fn=<SqueezeBackward1>)) |
| |
| >>> def adder(x, y): |
| ... return 2 * x + 3 * y |
| >>> inputs = (torch.rand(2), torch.rand(2)) |
| >>> v = (torch.ones(2), torch.ones(2)) |
| >>> jvp(adder, inputs, v) |
| (tensor([2.2399, 2.5005]), |
| tensor([5., 5.])) |
| |
| """ |
| |
| with torch.enable_grad(): |
| is_inputs_tuple, inputs = _as_tuple(inputs, "inputs", "jvp") |
| inputs = _grad_preprocess(inputs, create_graph=create_graph, need_graph=True) |
| |
| if v is not None: |
| _, v = _as_tuple(v, "v", "jvp") |
| v = _grad_preprocess(v, create_graph=create_graph, need_graph=False) |
| _validate_v(v, inputs, is_inputs_tuple) |
| else: |
| if len(inputs) != 1 or inputs[0].nelement() != 1: |
| raise RuntimeError( |
| "The vector v can only be None if the input to " |
| "the user-provided function is a single Tensor " |
| "with a single element." |
| ) |
| |
| outputs = func(*inputs) |
| is_outputs_tuple, outputs = _as_tuple( |
| outputs, "outputs of the user-provided function", "jvp" |
| ) |
| _check_requires_grad(outputs, "outputs", strict=strict) |
| # The backward is linear so the value of grad_outputs is not important as |
| # it won't appear in the double backward graph. We only need to ensure that |
| # it does not contain inf or nan. |
| grad_outputs = tuple( |
| torch.zeros_like(out, requires_grad=True) for out in outputs |
| ) |
| |
| grad_inputs = _autograd_grad(outputs, inputs, grad_outputs, create_graph=True) |
| _check_requires_grad(grad_inputs, "grad_inputs", strict=strict) |
| |
| if create_graph: |
| with torch.enable_grad(): |
| grad_res = _autograd_grad( |
| grad_inputs, grad_outputs, v, create_graph=create_graph |
| ) |
| jvp = _fill_in_zeros(grad_res, outputs, strict, create_graph, "back_trick") |
| else: |
| grad_res = _autograd_grad( |
| grad_inputs, grad_outputs, v, create_graph=create_graph |
| ) |
| jvp = _fill_in_zeros(grad_res, outputs, strict, create_graph, "back_trick") |
| |
| # Cleanup objects and return them to the user |
| outputs = _grad_postprocess(outputs, create_graph) |
| jvp = _grad_postprocess(jvp, create_graph) |
| |
| return _tuple_postprocess(outputs, is_outputs_tuple), _tuple_postprocess( |
| jvp, is_outputs_tuple |
| ) |
| |
| |
| def _construct_standard_basis_for( |
| tensors: Tuple[torch.Tensor, ...], tensor_numels: Tuple[int, ...] |
| ) -> Tuple[torch.Tensor, ...]: |
| # This function: |
| # - constructs a N=sum(tensor_numels) standard basis. i.e. an NxN identity matrix. |
| # - Splits the identity matrix into chunks with each chunk size determined by `tensor_numels`. |
| # - Each chunk corresponds to one tensor. The chunk has the same dtype and |
| # device as the tensor |
| # |
| # For example, with tensor_numels = [1, 2, 1], this function returns: |
| # ( tensor([[1], tensor([[0, 0], tensor([[0], |
| # [0], [1, 0], [0], |
| # [0], [0, 1], [0], |
| # [0]]) , [0, 0]]) , [1]]) ) |
| # |
| # Precondition: tensor_numels == tuple(tensor.numel() for tensor in tensors) |
| # Precondition: tensors always has at least one element. |
| # |
| # See NOTE: [Computing jacobian with vmap and grad for multiple tensors] |
| # for context behind this function. All the pre-conditions are guarded for |
| # in torch.autograd.functional.jacobian. |
| assert len(tensors) == len(tensor_numels) |
| assert len(tensors) > 0 |
| total_numel = sum(tensor_numels) |
| chunks = tuple( |
| tensor.new_zeros(total_numel, tensor_numel) |
| for tensor, tensor_numel in zip(tensors, tensor_numels) |
| ) |
| diag_start_idx = 0 |
| for chunk, numel in zip(chunks, tensor_numels): |
| chunk.diagonal(diag_start_idx).fill_(1) |
| diag_start_idx -= numel |
| return chunks |
| |
| |
| def _jacfwd(func, inputs, strict=False, vectorize=False): |
| if strict: |
| raise RuntimeError( |
| "torch.autograd.functional.jacobian: `strict=True` " |
| 'and `strategy="forward-mode"` are not supported together (yet). ' |
| "Please either set `strict=False` or " |
| '`strategy="reverse-mode"`.' |
| ) |
| is_inputs_tuple, inputs = _as_tuple(inputs, "inputs", "jacobian") |
| output_info = [] |
| |
| if vectorize: |
| # See NOTE: [Computing jacobian with vmap and grad for multiple outputs] |
| input_numels = tuple(input.numel() for input in inputs) |
| |
| # Step 1: Prepare tangents |
| tangents = _construct_standard_basis_for(inputs, input_numels) |
| |
| # Step 2: Compute vmap over computation with dual tensors |
| def jvp(tangents): |
| with fwAD.dual_level(): |
| dual_inputs = tuple( |
| fwAD.make_dual(input, tangent.view_as(input)) |
| for input, tangent in zip(inputs, tangents) |
| ) |
| _is_outputs_tuple, dual_outputs = _as_tuple( |
| func(*dual_inputs), "outputs" |
| ) |
| output_info.append(_is_outputs_tuple) |
| jv = [] |
| primal_outs = [] |
| for dual_out in dual_outputs: |
| primal, tangent = fwAD.unpack_dual(dual_out) |
| primal_outs.append(primal) |
| if tangent is not None: |
| jv.append(tangent) |
| else: |
| jv.append(torch.zeros_like(primal)) |
| output_info.append(primal_outs) |
| return tuple(jv) |
| |
| outputs_before_split = _vmap(jvp)(tangents) |
| is_outputs_tuple, outputs = output_info |
| # Step 3: for each of the output tangents, split along dim 0 |
| jacobian_input_output = [] |
| for jac, output_i in zip(outputs_before_split, outputs): |
| jacobian_output_i_output = [] |
| for jac, input_j in zip(jac.split(input_numels, dim=0), inputs): |
| # We need to transpose the Jacobian because in forward AD, the |
| # batch dimension represents that of the inputs |
| jacobian_input_i_output_j = jac.permute(*range(1, jac.ndim), 0).reshape( |
| (*output_i.shape, *input_j.shape) |
| ) # noqa: C409 |
| |
| jacobian_output_i_output.append(jacobian_input_i_output_j) |
| jacobian_input_output.append(jacobian_output_i_output) |
| |
| # Omit [Step 4] because everything is already transposed w/ forward AD |
| return _tuple_postprocess( |
| jacobian_input_output, (is_outputs_tuple, is_inputs_tuple) |
| ) |
| else: |
| raise NotImplementedError( |
| "Computing Jacobian using forward-AD or forward-over-reverse Hessian is" |
| "only implemented for `vectorize=True`." |
| ) |
| |
| |
| def jacobian( |
| func, |
| inputs, |
| create_graph=False, |
| strict=False, |
| vectorize=False, |
| strategy="reverse-mode", |
| ): |
| r"""Function that computes the Jacobian of a given function. |
| |
| Args: |
| func (function): a Python function that takes Tensor inputs and returns |
| a tuple of Tensors or a Tensor. |
| inputs (tuple of Tensors or Tensor): inputs to the function ``func``. |
| create_graph (bool, optional): If ``True``, the Jacobian will be |
| computed in a differentiable manner. Note that when ``strict`` is |
| ``False``, the result can not require gradients or be disconnected |
| from the inputs. Defaults to ``False``. |
| strict (bool, optional): If ``True``, an error will be raised when we |
| detect that there exists an input such that all the outputs are |
| independent of it. If ``False``, we return a Tensor of zeros as the |
| jacobian for said inputs, which is the expected mathematical value. |
| Defaults to ``False``. |
| vectorize (bool, optional): This feature is experimental. |
| Please consider using :func:`torch.func.jacrev` or |
| :func:`torch.func.jacfwd` instead if you are looking for something |
| less experimental and more performant. |
| When computing the jacobian, usually we invoke |
| ``autograd.grad`` once per row of the jacobian. If this flag is |
| ``True``, we perform only a single ``autograd.grad`` call with |
| ``batched_grad=True`` which uses the vmap prototype feature. |
| Though this should lead to performance improvements in many cases, |
| because this feature is still experimental, there may be performance |
| cliffs. See :func:`torch.autograd.grad`'s ``batched_grad`` parameter for |
| more information. |
| strategy (str, optional): Set to ``"forward-mode"`` or ``"reverse-mode"`` to |
| determine whether the Jacobian will be computed with forward or reverse |
| mode AD. Currently, ``"forward-mode"`` requires ``vectorized=True``. |
| Defaults to ``"reverse-mode"``. If ``func`` has more outputs than |
| inputs, ``"forward-mode"`` tends to be more performant. Otherwise, |
| prefer to use ``"reverse-mode"``. |
| |
| Returns: |
| Jacobian (Tensor or nested tuple of Tensors): if there is a single |
| input and output, this will be a single Tensor containing the |
| Jacobian for the linearized inputs and output. If one of the two is |
| a tuple, then the Jacobian will be a tuple of Tensors. If both of |
| them are tuples, then the Jacobian will be a tuple of tuple of |
| Tensors where ``Jacobian[i][j]`` will contain the Jacobian of the |
| ``i``\th output and ``j``\th input and will have as size the |
| concatenation of the sizes of the corresponding output and the |
| corresponding input and will have same dtype and device as the |
| corresponding input. If strategy is ``forward-mode``, the dtype will be |
| that of the output; otherwise, the input. |
| |
| Example: |
| |
| >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_AUTOGRAD) |
| >>> def exp_reducer(x): |
| ... return x.exp().sum(dim=1) |
| >>> inputs = torch.rand(2, 2) |
| >>> # xdoctest: +IGNORE_WANT("non-deterministic") |
| >>> jacobian(exp_reducer, inputs) |
| tensor([[[1.4917, 2.4352], |
| [0.0000, 0.0000]], |
| [[0.0000, 0.0000], |
| [2.4369, 2.3799]]]) |
| |
| >>> jacobian(exp_reducer, inputs, create_graph=True) |
| tensor([[[1.4917, 2.4352], |
| [0.0000, 0.0000]], |
| [[0.0000, 0.0000], |
| [2.4369, 2.3799]]], grad_fn=<ViewBackward>) |
| |
| >>> def exp_adder(x, y): |
| ... return 2 * x.exp() + 3 * y |
| >>> inputs = (torch.rand(2), torch.rand(2)) |
| >>> jacobian(exp_adder, inputs) |
| (tensor([[2.8052, 0.0000], |
| [0.0000, 3.3963]]), |
| tensor([[3., 0.], |
| [0., 3.]])) |
| """ |
| assert strategy in ("forward-mode", "reverse-mode"), ( |
| 'Expected strategy to be either "forward-mode" or "reverse-mode". Hint: If your ' |
| 'function has more outputs than inputs, "forward-mode" tends to be more performant. ' |
| 'Otherwise, prefer to use "reverse-mode".' |
| ) |
| if strategy == "forward-mode": |
| if create_graph: |
| raise NotImplementedError( |
| "torch.autograd.functional.jacobian: `create_graph=True` " |
| 'and `strategy="forward-mode"` are not supported together (yet). ' |
| "Please either set `create_graph=False` or " |
| '`strategy="reverse-mode"`.' |
| ) |
| return _jacfwd(func, inputs, strict, vectorize) |
| |
| with torch.enable_grad(): |
| is_inputs_tuple, inputs = _as_tuple(inputs, "inputs", "jacobian") |
| inputs = _grad_preprocess(inputs, create_graph=create_graph, need_graph=True) |
| |
| outputs = func(*inputs) |
| is_outputs_tuple, outputs = _as_tuple( |
| outputs, "outputs of the user-provided function", "jacobian" |
| ) |
| _check_requires_grad(outputs, "outputs", strict=strict) |
| |
| if vectorize: |
| if strict: |
| raise RuntimeError( |
| "torch.autograd.functional.jacobian: `strict=True` " |
| "and `vectorized=True` are not supported together. " |
| "Please either set `strict=False` or " |
| "`vectorize=False`." |
| ) |
| # NOTE: [Computing jacobian with vmap and grad for multiple outputs] |
| # |
| # Let's consider f(x) = (x**2, x.sum()) and let x = torch.randn(3). |
| # It turns out we can compute the jacobian of this function with a single |
| # call to autograd.grad by using vmap over the correct grad_outputs. |
| # |
| # Firstly, one way to compute the jacobian is to stack x**2 and x.sum() |
| # into a 4D vector. E.g., use g(x) = torch.stack([x**2, x.sum()]) |
| # |
| # To get the first row of the jacobian, we call |
| # >>> autograd.grad(g(x), x, grad_outputs=torch.tensor([1, 0, 0, 0])) |
| # To get the 2nd row of the jacobian, we call |
| # >>> autograd.grad(g(x), x, grad_outputs=torch.tensor([0, 1, 0, 0])) |
| # and so on. |
| # |
| # Using vmap, we can vectorize all 4 of these computations into one by |
| # passing the standard basis for R^4 as the grad_output. |
| # vmap(partial(autograd.grad, g(x), x))(torch.eye(4)). |
| # |
| # Now, how do we compute the jacobian *without stacking the output*? |
| # We can just split the standard basis across the outputs. So to |
| # compute the jacobian of f(x), we'd use |
| # >>> autograd.grad(f(x), x, grad_outputs=_construct_standard_basis_for(...)) |
| # The grad_outputs looks like the following: |
| # ( torch.tensor([[1, 0, 0], |
| # [0, 1, 0], |
| # [0, 0, 1], |
| # [0, 0, 0]]), |
| # torch.tensor([[0], |
| # [0], |
| # [0], |
| # [1]]) ) |
| # |
| # But we're not done yet! |
| # >>> vmap(partial(autograd.grad(f(x), x, grad_outputs=...))) |
| # returns a Tensor of shape [4, 3]. We have to remember to split the |
| # jacobian of shape [4, 3] into two: |
| # - one of shape [3, 3] for the first output |
| # - one of shape [ 3] for the second output |
| |
| # Step 1: Construct grad_outputs by splitting the standard basis |
| output_numels = tuple(output.numel() for output in outputs) |
| grad_outputs = _construct_standard_basis_for(outputs, output_numels) |
| flat_outputs = tuple(output.reshape(-1) for output in outputs) |
| |
| # Step 2: Call vmap + autograd.grad |
| def vjp(grad_output): |
| vj = list( |
| _autograd_grad( |
| flat_outputs, |
| inputs, |
| grad_output, |
| create_graph=create_graph, |
| is_grads_batched=True, |
| ) |
| ) |
| for el_idx, vj_el in enumerate(vj): |
| if vj_el is not None: |
| continue |
| vj[el_idx] = torch.zeros_like(inputs[el_idx]).expand( |
| (sum(output_numels),) + inputs[el_idx].shape |
| ) |
| return tuple(vj) |
| |
| jacobians_of_flat_output = vjp(grad_outputs) |
| |
| # Step 3: The returned jacobian is one big tensor per input. In this step, |
| # we split each Tensor by output. |
| jacobian_input_output = [] |
| for jac, input_i in zip(jacobians_of_flat_output, inputs): |
| jacobian_input_i_output = [] |
| for jac, output_j in zip(jac.split(output_numels, dim=0), outputs): |
| jacobian_input_i_output_j = jac.view(output_j.shape + input_i.shape) |
| jacobian_input_i_output.append(jacobian_input_i_output_j) |
| jacobian_input_output.append(jacobian_input_i_output) |
| |
| # Step 4: Right now, `jacobian` is a List[List[Tensor]]. |
| # The outer List corresponds to the number of inputs, |
| # the inner List corresponds to the number of outputs. |
| # We need to exchange the order of these and convert to tuples |
| # before returning. |
| jacobian_output_input = tuple(zip(*jacobian_input_output)) |
| |
| jacobian_output_input = _grad_postprocess( |
| jacobian_output_input, create_graph |
| ) |
| return _tuple_postprocess( |
| jacobian_output_input, (is_outputs_tuple, is_inputs_tuple) |
| ) |
| |
| jacobian: Tuple[torch.Tensor, ...] = tuple() |
| |
| for i, out in enumerate(outputs): |
| # mypy complains that expression and variable have different types due to the empty list |
| jac_i: Tuple[List[torch.Tensor]] = tuple([] for _ in range(len(inputs))) # type: ignore[assignment] |
| for j in range(out.nelement()): |
| vj = _autograd_grad( |
| (out.reshape(-1)[j],), |
| inputs, |
| retain_graph=True, |
| create_graph=create_graph, |
| ) |
| |
| for el_idx, (jac_i_el, vj_el, inp_el) in enumerate( |
| zip(jac_i, vj, inputs) |
| ): |
| if vj_el is not None: |
| if strict and create_graph and not vj_el.requires_grad: |
| msg = ( |
| "The jacobian of the user-provided function is " |
| f"independent of input {i}. This is not allowed in " |
| "strict mode when create_graph=True." |
| ) |
| raise RuntimeError(msg) |
| jac_i_el.append(vj_el) |
| else: |
| if strict: |
| msg = ( |
| f"Output {i} of the user-provided function is " |
| f"independent of input {el_idx}. This is not allowed in " |
| "strict mode." |
| ) |
| raise RuntimeError(msg) |
| jac_i_el.append(torch.zeros_like(inp_el)) |
| |
| jacobian += ( |
| tuple( |
| torch.stack(jac_i_el, dim=0).view( |
| out.size() + inputs[el_idx].size() # type: ignore[operator] |
| ) |
| for (el_idx, jac_i_el) in enumerate(jac_i) |
| ), |
| ) |
| |
| jacobian = _grad_postprocess(jacobian, create_graph) |
| |
| return _tuple_postprocess(jacobian, (is_outputs_tuple, is_inputs_tuple)) |
| |
| |
| def hessian( |
| func, |
| inputs, |
| create_graph=False, |
| strict=False, |
| vectorize=False, |
| outer_jacobian_strategy="reverse-mode", |
| ): |
| r"""Function that computes the Hessian of a given scalar function. |
| |
| Args: |
| func (function): a Python function that takes Tensor inputs and returns |
| a Tensor with a single element. |
| inputs (tuple of Tensors or Tensor): inputs to the function ``func``. |
| create_graph (bool, optional): If ``True``, the Hessian will be computed in |
| a differentiable manner. Note that when ``strict`` is ``False``, the result can not |
| require gradients or be disconnected from the inputs. |
| Defaults to ``False``. |
| strict (bool, optional): If ``True``, an error will be raised when we detect that there exists an input |
| such that all the outputs are independent of it. If ``False``, we return a Tensor of zeros as the |
| hessian for said inputs, which is the expected mathematical value. |
| Defaults to ``False``. |
| vectorize (bool, optional): This feature is experimental. |
| Please consider using :func:`torch.func.hessian` |
| instead if you are looking for something less experimental and more performant. |
| When computing the hessian, usually we invoke |
| ``autograd.grad`` once per row of the hessian. If this flag is |
| ``True``, we use the vmap prototype feature as the backend to |
| vectorize calls to ``autograd.grad`` so we only invoke it once |
| instead of once per row. This should lead to performance |
| improvements in many use cases, however, due to this feature |
| being incomplete, there may be performance cliffs. Please |
| use `torch._C._debug_only_display_vmap_fallback_warnings(True)` |
| to show any performance warnings and file us issues if |
| warnings exist for your use case. Defaults to ``False``. |
| outer_jacobian_strategy (str, optional): The Hessian is computed by |
| computing the Jacobian of a Jacobian. The inner Jacobian is always |
| computed in reverse-mode AD. Setting strategy to ``"forward-mode"`` |
| or ``"reverse-mode"`` determines whether the outer Jacobian will be |
| computed with forward or reverse mode AD. Currently, computing the outer |
| Jacobian in ``"forward-mode"`` requires ``vectorized=True``. Defaults |
| to ``"reverse-mode"``. |
| |
| Returns: |
| Hessian (Tensor or a tuple of tuple of Tensors): if there is a single input, |
| this will be a single Tensor containing the Hessian for the input. |
| If it is a tuple, then the Hessian will be a tuple of tuples where |
| ``Hessian[i][j]`` will contain the Hessian of the ``i``\th input |
| and ``j``\th input with size the sum of the size of the ``i``\th input plus |
| the size of the ``j``\th input. ``Hessian[i][j]`` will have the same |
| dtype and device as the corresponding ``i``\th input. |
| |
| Example: |
| |
| >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_AUTOGRAD) |
| >>> def pow_reducer(x): |
| ... return x.pow(3).sum() |
| >>> inputs = torch.rand(2, 2) |
| >>> # xdoctest: +IGNORE_WANT("non-deterministic") |
| >>> hessian(pow_reducer, inputs) |
| tensor([[[[5.2265, 0.0000], |
| [0.0000, 0.0000]], |
| [[0.0000, 4.8221], |
| [0.0000, 0.0000]]], |
| [[[0.0000, 0.0000], |
| [1.9456, 0.0000]], |
| [[0.0000, 0.0000], |
| [0.0000, 3.2550]]]]) |
| |
| >>> hessian(pow_reducer, inputs, create_graph=True) |
| tensor([[[[5.2265, 0.0000], |
| [0.0000, 0.0000]], |
| [[0.0000, 4.8221], |
| [0.0000, 0.0000]]], |
| [[[0.0000, 0.0000], |
| [1.9456, 0.0000]], |
| [[0.0000, 0.0000], |
| [0.0000, 3.2550]]]], grad_fn=<ViewBackward>) |
| |
| |
| >>> def pow_adder_reducer(x, y): |
| ... return (2 * x.pow(2) + 3 * y.pow(2)).sum() |
| >>> inputs = (torch.rand(2), torch.rand(2)) |
| >>> hessian(pow_adder_reducer, inputs) |
| ((tensor([[4., 0.], |
| [0., 4.]]), |
| tensor([[0., 0.], |
| [0., 0.]])), |
| (tensor([[0., 0.], |
| [0., 0.]]), |
| tensor([[6., 0.], |
| [0., 6.]]))) |
| """ |
| |
| is_inputs_tuple, inputs = _as_tuple(inputs, "inputs", "hessian") |
| assert outer_jacobian_strategy in ( |
| "forward-mode", |
| "reverse-mode", |
| ), 'Expected strategy to be either "forward-mode" or "reverse-mode".' |
| |
| def ensure_single_output_function(*inp): |
| out = func(*inp) |
| is_out_tuple, t_out = _as_tuple( |
| out, "outputs of the user-provided function", "hessian" |
| ) |
| _check_requires_grad(t_out, "outputs", strict=strict) |
| |
| if is_out_tuple or not isinstance(out, torch.Tensor): |
| raise RuntimeError( |
| "The function given to hessian should return a single Tensor" |
| ) |
| |
| if out.nelement() != 1: |
| raise RuntimeError( |
| "The Tensor returned by the function given to hessian should contain a single element" |
| ) |
| |
| return out.squeeze() |
| |
| def jac_func(*inp): |
| if outer_jacobian_strategy == "forward-mode": |
| # _grad_preprocess requires create_graph=True and input to require_grad |
| # or else the input will be detached |
| inp = tuple(t.requires_grad_(True) for t in inp) |
| jac = jacobian(ensure_single_output_function, inp, create_graph=True) |
| _check_requires_grad(jac, "jacobian", strict=strict) |
| return jac |
| |
| res = jacobian( |
| jac_func, |
| inputs, |
| create_graph=create_graph, |
| strict=strict, |
| vectorize=vectorize, |
| strategy=outer_jacobian_strategy, |
| ) |
| return _tuple_postprocess(res, (is_inputs_tuple, is_inputs_tuple)) |
| |
| |
| def vhp(func, inputs, v=None, create_graph=False, strict=False): |
| r"""Function that computes the dot product between a vector ``v`` and the |
| Hessian of a given scalar function at the point given by the inputs. |
| |
| Args: |
| func (function): a Python function that takes Tensor inputs and returns |
| a Tensor with a single element. |
| inputs (tuple of Tensors or Tensor): inputs to the function ``func``. |
| v (tuple of Tensors or Tensor): The vector for which the vector Hessian |
| product is computed. Must be the same size as the input of |
| ``func``. This argument is optional when ``func``'s input contains |
| a single element and (if it is not provided) will be set as a |
| Tensor containing a single ``1``. |
| create_graph (bool, optional): If ``True``, both the output and result |
| will be computed in a differentiable way. Note that when ``strict`` |
| is ``False``, the result can not require gradients or be |
| disconnected from the inputs. |
| Defaults to ``False``. |
| strict (bool, optional): If ``True``, an error will be raised when we |
| detect that there exists an input such that all the outputs are |
| independent of it. If ``False``, we return a Tensor of zeros as the |
| vhp for said inputs, which is the expected mathematical value. |
| Defaults to ``False``. |
| |
| Returns: |
| output (tuple): tuple with: |
| func_output (tuple of Tensors or Tensor): output of ``func(inputs)`` |
| |
| vhp (tuple of Tensors or Tensor): result of the dot product with the |
| same shape as the inputs. |
| |
| Example: |
| |
| >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_AUTOGRAD) |
| >>> def pow_reducer(x): |
| ... return x.pow(3).sum() |
| >>> inputs = torch.rand(2, 2) |
| >>> v = torch.ones(2, 2) |
| >>> # xdoctest: +IGNORE_WANT("non-deterministic") |
| >>> vhp(pow_reducer, inputs, v) |
| (tensor(0.5591), |
| tensor([[1.0689, 1.2431], |
| [3.0989, 4.4456]])) |
| >>> vhp(pow_reducer, inputs, v, create_graph=True) |
| (tensor(0.5591, grad_fn=<SumBackward0>), |
| tensor([[1.0689, 1.2431], |
| [3.0989, 4.4456]], grad_fn=<MulBackward0>)) |
| >>> def pow_adder_reducer(x, y): |
| ... return (2 * x.pow(2) + 3 * y.pow(2)).sum() |
| >>> inputs = (torch.rand(2), torch.rand(2)) |
| >>> v = (torch.zeros(2), torch.ones(2)) |
| >>> vhp(pow_adder_reducer, inputs, v) |
| (tensor(4.8053), |
| (tensor([0., 0.]), |
| tensor([6., 6.]))) |
| """ |
| |
| with torch.enable_grad(): |
| is_inputs_tuple, inputs = _as_tuple(inputs, "inputs", "vhp") |
| inputs = _grad_preprocess(inputs, create_graph=create_graph, need_graph=True) |
| |
| if v is not None: |
| _, v = _as_tuple(v, "v", "vhp") |
| v = _grad_preprocess(v, create_graph=create_graph, need_graph=False) |
| _validate_v(v, inputs, is_inputs_tuple) |
| else: |
| if len(inputs) != 1 or inputs[0].nelement() != 1: |
| raise RuntimeError( |
| "The vector v can only be None if the input to the user-provided function " |
| "is a single Tensor with a single element." |
| ) |
| outputs = func(*inputs) |
| is_outputs_tuple, outputs = _as_tuple( |
| outputs, "outputs of the user-provided function", "vhp" |
| ) |
| _check_requires_grad(outputs, "outputs", strict=strict) |
| |
| if is_outputs_tuple or not isinstance(outputs[0], torch.Tensor): |
| raise RuntimeError( |
| "The function given to vhp should return a single Tensor" |
| ) |
| |
| if outputs[0].nelement() != 1: |
| raise RuntimeError( |
| "The Tensor returned by the function given to vhp should contain a single element" |
| ) |
| |
| jac = _autograd_grad(outputs, inputs, create_graph=True) |
| _check_requires_grad(jac, "jacobian", strict=strict) |
| |
| enable_grad = True if create_graph else torch.is_grad_enabled() |
| with torch.set_grad_enabled(enable_grad): |
| grad_res = _autograd_grad(jac, inputs, v, create_graph=create_graph) |
| vhp = _fill_in_zeros(grad_res, inputs, strict, create_graph, "double_back") |
| |
| outputs = _grad_postprocess(outputs, create_graph) |
| vhp = _grad_postprocess(vhp, create_graph) |
| |
| return _tuple_postprocess(outputs, is_outputs_tuple), _tuple_postprocess( |
| vhp, is_inputs_tuple |
| ) |
| |
| |
| def hvp(func, inputs, v=None, create_graph=False, strict=False): |
| r"""Function that computes the dot product between the Hessian of a given scalar |
| function and a vector ``v`` at the point given by the inputs. |
| |
| Args: |
| func (function): a Python function that takes Tensor inputs and returns |
| a Tensor with a single element. |
| inputs (tuple of Tensors or Tensor): inputs to the function ``func``. |
| v (tuple of Tensors or Tensor): The vector for which the Hessian vector |
| product is computed. Must be the same size as the input of |
| ``func``. This argument is optional when ``func``'s input contains |
| a single element and (if it is not provided) will be set as a |
| Tensor containing a single ``1``. |
| create_graph (bool, optional): If ``True``, both the output and result will be |
| computed in a differentiable way. Note that when ``strict`` is |
| ``False``, the result can not require gradients or be disconnected |
| from the inputs. Defaults to ``False``. |
| strict (bool, optional): If ``True``, an error will be raised when we |
| detect that there exists an input such that all the outputs are |
| independent of it. If ``False``, we return a Tensor of zeros as the |
| hvp for said inputs, which is the expected mathematical value. |
| Defaults to ``False``. |
| Returns: |
| output (tuple): tuple with: |
| func_output (tuple of Tensors or Tensor): output of ``func(inputs)`` |
| |
| hvp (tuple of Tensors or Tensor): result of the dot product with |
| the same shape as the inputs. |
| |
| Example: |
| |
| >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_AUTOGRAD) |
| >>> def pow_reducer(x): |
| ... return x.pow(3).sum() |
| >>> inputs = torch.rand(2, 2) |
| >>> v = torch.ones(2, 2) |
| >>> # xdoctest: +IGNORE_WANT("non-deterministic") |
| >>> hvp(pow_reducer, inputs, v) |
| (tensor(0.1448), |
| tensor([[2.0239, 1.6456], |
| [2.4988, 1.4310]])) |
| |
| >>> hvp(pow_reducer, inputs, v, create_graph=True) |
| (tensor(0.1448, grad_fn=<SumBackward0>), |
| tensor([[2.0239, 1.6456], |
| [2.4988, 1.4310]], grad_fn=<MulBackward0>)) |
| |
| |
| >>> def pow_adder_reducer(x, y): |
| ... return (2 * x.pow(2) + 3 * y.pow(2)).sum() |
| >>> inputs = (torch.rand(2), torch.rand(2)) |
| >>> v = (torch.zeros(2), torch.ones(2)) |
| >>> hvp(pow_adder_reducer, inputs, v) |
| (tensor(2.3030), |
| (tensor([0., 0.]), |
| tensor([6., 6.]))) |
| |
| Note: |
| |
| This function is significantly slower than `vhp` due to backward mode AD constraints. |
| If your functions is twice continuously differentiable, then hvp = vhp.t(). So if you |
| know that your function satisfies this condition, you should use vhp instead that is |
| much faster with the current implementation. |
| |
| """ |
| |
| with torch.enable_grad(): |
| is_inputs_tuple, inputs = _as_tuple(inputs, "inputs", "hvp") |
| inputs = _grad_preprocess(inputs, create_graph=create_graph, need_graph=True) |
| |
| if v is not None: |
| _, v = _as_tuple(v, "v", "hvp") |
| v = _grad_preprocess(v, create_graph=create_graph, need_graph=False) |
| _validate_v(v, inputs, is_inputs_tuple) |
| else: |
| if len(inputs) != 1 or inputs[0].nelement() != 1: |
| raise RuntimeError( |
| "The vector v can only be None if the input to the user-provided function " |
| "is a single Tensor with a single element." |
| ) |
| outputs = func(*inputs) |
| is_outputs_tuple, outputs = _as_tuple( |
| outputs, "outputs of the user-provided function", "hvp" |
| ) |
| _check_requires_grad(outputs, "outputs", strict=strict) |
| |
| if is_outputs_tuple or not isinstance(outputs[0], torch.Tensor): |
| raise RuntimeError( |
| "The function given to hvp should return a single Tensor" |
| ) |
| |
| if outputs[0].nelement() != 1: |
| raise RuntimeError( |
| "The Tensor returned by the function given to hvp should contain a single element" |
| ) |
| |
| jac = _autograd_grad(outputs, inputs, create_graph=True) |
| _check_requires_grad(jac, "jacobian", strict=strict) |
| |
| grad_jac = tuple(torch.zeros_like(inp, requires_grad=True) for inp in inputs) |
| |
| double_back = _autograd_grad(jac, inputs, grad_jac, create_graph=True) |
| _check_requires_grad(jac, "hessian", strict=strict) |
| |
| enable_grad = True if create_graph else torch.is_grad_enabled() |
| with torch.set_grad_enabled(enable_grad): |
| grad_res = _autograd_grad(double_back, grad_jac, v, create_graph=create_graph) |
| hvp = _fill_in_zeros( |
| grad_res, inputs, strict, create_graph, "double_back_trick" |
| ) |
| |
| outputs = _grad_postprocess(outputs, create_graph) |
| hvp = _grad_postprocess(hvp, create_graph) |
| |
| return _tuple_postprocess(outputs, is_outputs_tuple), _tuple_postprocess( |
| hvp, is_inputs_tuple |
| ) |
