Repository: pytorch
Issue: Functional Jacobian does not work with Torchdiffeq
Labels: module: autograd,triaged

### 🐛 Describe the bug

I am following the implementation of #49171 to obtain the numerical jacobian wrt model parameters. My model, however, is a Neural ODE defined following torchdiffeq. I observe that the jacobian returns all zeros when I wrap my neural network in an ODESolve. Please kindly point out whether this is a feature of `autograd.jacobian`, or perhaps I misunderstood something. Thanks a lot!

```
import time
import os
import argparse
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.optim as optim
from torch.autograd.functional import jacobian
from torch.nn.utils import _stateless
from torchdiffeq import odeint_adjoint as odeint


class ODEfunc(nn.Module):

    def __init__(self, dim, nhidden):
        super(ODEfunc, self).__init__()
        self.elu = nn.ELU(inplace=True)
        self.fc1 = nn.Linear(dim, nhidden)
        self.fc2 = nn.Linear(nhidden, nhidden)
        self.fc3 = nn.Linear(nhidden, dim)
        self.nfe = 0

    def forward(self, t, x):
        self.nfe += 1
        out = self.fc1(x)
        out = self.elu(out)
        out = self.fc2(out)
        out = self.elu(out)
        out = self.fc3(out)
        return out


class ODEBlock(nn.Module):

    def __init__(self, odefunc, t):
        super(ODEBlock, self).__init__()
        self.odefunc = odefunc
        if len(t) == 2:
            self.integration_times = torch.tensor([t[0], t[1]]).float()
        else:
            self.integration_times = t.float()

    def forward(self, x):
        out = odeint(self.odefunc, x, self.integration_times, rtol=1e-3, atol=1e-3)
        return out

model = ODEBlock(ODEfunc(1, 20), torch.linspace(0, 10, 101)).to(device)
P = 0.01
K = 1
tspan = torch.linspace(0, 10, 101)
func = lambda x: K * P * torch.exp(x) / (K + P * torch.exp(x) - P)
logistic_data = []
for i in tspan:
    logistic_data.append([func(i).item()])
z0 = torch.tensor(logistic_data[:-1]).float().to(device)
def functional_loss(*functional_params):
    out: torch.Tensor = _stateless.functional_call(model, {n: p for n, p in zip(names, functional_params)}, z0)
    return out
names = list(n for n, _ in model.named_parameters())
temp = jacobian(functional_loss, tuple(model.parameters()))
print(temp)

### Versions

Collecting environment information...
PyTorch version: 1.11.0
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 12.3.1 (x86_64)
GCC version: Could not collect
Clang version: 13.1.6 (clang-1316.0.21.2.5)
CMake version: version 3.17.1
Libc version: N/A

Python version: 3.8.12 (default, Oct 12 2021, 06:23:56)  [Clang 10.0.0 ] (64-bit runtime)
Python platform: macOS-10.16-x86_64-i386-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
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

Versions of relevant libraries:
[pip3] backpack-for-pytorch==1.5.0
[pip3] functorch==0.1.0
[pip3] numpy==1.22.3
[pip3] pytorch-ranger==0.1.1
[pip3] torch==1.11.0
[pip3] torch-optimizer==0.1.0
[pip3] torchaudio==0.11.0
[pip3] torchdiffeq==0.2.2
[pip3] torchmetrics==0.7.3
[pip3] torchvision==0.12.0
[conda] backpack-for-pytorch      1.5.0                    pypi_0    pypi
[conda] blas                      1.0                         mkl    defaults
[conda] ffmpeg                    4.3                  h0a44026_0    pytorch
[conda] functorch                 0.1.0                    pypi_0    pypi
[conda] mkl                       2021.4.0           hecd8cb5_637    defaults
[conda] mkl-service               2.4.0            py38h9ed2024_0    defaults
[conda] mkl_fft                   1.3.1            py38h4ab4a9b_0    defaults
[conda] mkl_random                1.2.2            py38hb2f4e1b_0    defaults
[conda] numpy                     1.22.3                   pypi_0    pypi
[conda] numpy-base                1.21.2           py38he0bd621_0    defaults
[conda] pytorch                   1.11.0                  py3.8_0    pytorch
[conda] pytorch-ranger            0.1.1                    pypi_0    pypi
[conda] torch-optimizer           0.1.0                    pypi_0    pypi
[conda] torchaudio                0.11.0                 py38_cpu    pytorch
[conda] torchdiffeq               0.2.2                    pypi_0    pypi
[conda] torchmetrics              0.7.3                    pypi_0    pypi
[conda] torchvision               0.12.0                 py38_cpu    pytorch

cc @ezyang @albanD @zou3519 @gqchen @pearu @nikitaved @soulitzer @Lezcano @Varal7