Repository: pytorch
Issue: RuntimeError for hessian vector product with jvp
Labels: triaged,module: functorch

### 🐛 Describe the bug

Following this [example](https://pytorch.org/functorch/2.0/notebooks/jacobians_hessians.html) for functorch, I would like to implement hessian vector products using func.torch. When trying to use torch.func.jvp, I get a RuntimeError due to some immutability. This is a repost of a [topic](https://discuss.pytorch.org/t/hessian-vector-products-with-torch-func/184914) in the discussion forum, due to no response.

Minimal example:
```python
import torch

dimensions = (3, 2)
num_data = 10

x = torch.rand(num_data, dimensions[0])
y = torch.rand(num_data, dimensions[1])

model = torch.nn.Linear(*dimensions, bias=True)
vec = {name: torch.rand(*p.shape) for name, p in model.named_parameters() if p.requires_grad}
params = dict(model.named_parameters())


def model_func(new_params):
    outputs = torch.func.functional_call(model, new_params, (x,))
    return torch.nn.functional.mse_loss(outputs, y)


def hvp_revrev(f, primals, tangents):
    _, vjp_fn = torch.func.vjp(torch.func.grad(f), *primals)
    return vjp_fn(*tangents)[0]


# this works
hess_rev_vp = hvp_revrev(model_func, (params, ), (vec, ))

# raises error
hess_vp = torch.func.jvp(torch.func.grad(model_func), (params, ), (vec,))
```

Error message:
```
Traceback (most recent call last):
  File "/opt/homebrew/Caskroom/mambaforge/base/envs/pydvl/lib/python3.8/site-packages/IPython/core/interactiveshell.py", line 3508, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-2-959bc9ee399f>", line 1, in <module>
    runfile('/Users/kristof/Projects/pyDVL/default_console.py', wdir='/Users/kristof/Projects/pyDVL')
  File "/Users/kristof/Library/Application Support/JetBrains/IntelliJIdea2023.2/plugins/python/helpers/pydev/_pydev_bundle/pydev_umd.py", line 197, in runfile
    pydev_imports.execfile(filename, global_vars, local_vars)  # execute the script
  File "/Users/kristof/Library/Application Support/JetBrains/IntelliJIdea2023.2/plugins/python/helpers/pydev/_pydev_imps/_pydev_execfile.py", line 18, in execfile
    exec(compile(contents+"\n", file, 'exec'), glob, loc)
  File "/Users/kristof/Projects/pyDVL/default_console.py", line 246, in <module>
    hess_vp = torch.func.jvp(torch.func.grad(model_func), (params, ), (vec,))
  File "/opt/homebrew/Caskroom/mambaforge/base/envs/pydvl/lib/python3.8/site-packages/torch/_functorch/eager_transforms.py", line 919, in jvp
    return _jvp_with_argnums(func, primals, tangents, argnums=None, strict=strict, has_aux=has_aux)
  File "/opt/homebrew/Caskroom/mambaforge/base/envs/pydvl/lib/python3.8/site-packages/torch/_functorch/vmap.py", line 38, in fn
    return f(*args, **kwargs)
  File "/opt/homebrew/Caskroom/mambaforge/base/envs/pydvl/lib/python3.8/site-packages/torch/_functorch/eager_transforms.py", line 968, in _jvp_with_argnums
    result_duals = func(*duals)
  File "/opt/homebrew/Caskroom/mambaforge/base/envs/pydvl/lib/python3.8/site-packages/torch/_functorch/apis.py", line 363, in wrapper
    return eager_transforms.grad_impl(func, argnums, has_aux, args, kwargs)
  File "/opt/homebrew/Caskroom/mambaforge/base/envs/pydvl/lib/python3.8/site-packages/torch/_functorch/eager_transforms.py", line 1287, in grad_impl
    results = grad_and_value(func, argnums, has_aux=has_aux)(*args, **kwargs)
  File "/opt/homebrew/Caskroom/mambaforge/base/envs/pydvl/lib/python3.8/site-packages/torch/_functorch/vmap.py", line 38, in fn
    return f(*args, **kwargs)
  File "/opt/homebrew/Caskroom/mambaforge/base/envs/pydvl/lib/python3.8/site-packages/torch/_functorch/eager_transforms.py", line 1270, in wrapper
    flat_grad_input = _autograd_grad(flat_outputs, flat_diff_args, create_graph=True)
  File "/opt/homebrew/Caskroom/mambaforge/base/envs/pydvl/lib/python3.8/site-packages/torch/_functorch/eager_transforms.py", line 116, in _autograd_grad
    grad_inputs = torch.autograd.grad(diff_outputs, inputs, grad_outputs,
  File "/opt/homebrew/Caskroom/mambaforge/base/envs/pydvl/lib/python3.8/site-packages/torch/autograd/__init__.py", line 394, in grad
    result = Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
RuntimeError: ZeroTensors are immutable. Please use the materialized zero tensor obtained using .clone() if you want a mutable tensor.
```



### Versions

PyTorch version: 2.1.0
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 14.0 (arm64)
GCC version: Could not collect
Clang version: 15.0.0 (clang-1500.0.40.1)
CMake version: version 3.27.7
Libc version: N/A

Python version: 3.8.17 | packaged by conda-forge | (default, Jun 16 2023, 07:11:32)  [Clang 14.0.6 ] (64-bit runtime)
Python platform: macOS-14.0-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Apple M2 Max

Versions of relevant libraries:
[pip3] mypy==0.982
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.24.4
[pip3] pytorch-lightning==2.0.9
[pip3] torch==2.1.0
[pip3] torchmetrics==1.2.0
[pip3] torchvision==0.16.0
[conda] numpy                     1.24.4                   pypi_0    pypi
[conda] pytorch-lightning         2.0.9                    pypi_0    pypi
[conda] torch                     2.1.0                    pypi_0    pypi
[conda] torchmetrics              1.2.0                    pypi_0    pypi
[conda] torchvision               0.16.0                   pypi_0    pypi

cc @zou3519 @Chillee @samdow @kshitij12345 @janeyx99