I have written some code which allows the user to create vector-valued functions of the form `f(x, y, theta)`

, where `theta`

is a dict containing parameters. The idea is that `x`

and `y`

will always be a part of the function but `theta`

could be any number of different parameters depending on the function of interest. Note that I am not using these functions as part of a neural network module or anything like that. I am purely using the autograd functionality to differentiate functions.

For simplicity in this example, here is a scalar input-output function:

```
def f(x, y, theta):
a, b = theta["a"], theta["b"]
return x**4 + y**3 + a*b**2
```

Then I’d like to take the Jacobian as in

```
import torch
from torch.autograd.functional import jacobian
x = torch.tensor(5.)
y = torch.tensor(3.)
theta = {"a": 2, "b": 6}
jacobian(f, (x, y, theta))
```

This of course does not work since the `jacobian`

expects a tuple of tensors as input. The error says that the function got a dict instead of a tensor or tuple of tensors. I tried to cast the dict to a tensor but this is not possible. I tried looking into `Tensordict`

but this does not appear to be accepted by the `jacobian`

function. I also tried passing kwargs into the function which get unpacked in the function body. This does not work because `jacobian`

is expecting `f()`

to take two positional arguments (just `x`

and `y`

).

In Jax it is possible to differentiate w.r.t. a dict input and I was originally able to obtain the correct answer in that framework. Is there any way I could do this in PyTorch?