Come faccio a utilizzare autograd per una funzione separata indipendente di backpropagate in PyTorch?

Ho due variabili, x e theta. Sto cercando di ridurre al minimo la mia perdita, rispetto al theta solo, ma come parte della mia perdita di funzione di cui ho bisogno la derivata di una funzione diversa (f) rispetto al x. Questo derivato stesso non è rilevante per la riduzione al minimo, solo la sua uscita. Tuttavia, quando l'attuazione di questo PyTorch ricevo un errore di Runtime.

Un esempio minimale è come segue:

# minimal example of two different autograds
import torch

from torch.autograd.functional import jacobian
def f(theta, x):
    return torch.sum(theta * x ** 2)

def df(theta, x):
    J = jacobian(lambda x: f(theta, x), x)
    return J

# example evaluations of the autograd gradient
x = torch.tensor([1., 2.])
theta = torch.tensor([1., 1.], requires_grad = True)

# derivative should be 2*theta*x (same as an analytical)
with torch.no_grad():
    print(df(theta, x))
    print(2*theta*x)

tensore([2., 4.])

tensore([2., 4.])

# define some arbitrary loss as a fn of theta
loss = torch.sum(df(theta, x)**2)
loss.backward()

dà il seguente errore

RuntimeError: l'elemento 0 di tensori non richiede grad e non ha un grad_fn

Se ho un'analitica derivati (2*theta*x), funziona benissimo:

loss = torch.sum((2*theta*x)**2)
loss.backward()

C'è un modo per fare questo in PyTorch? O sono io limitato in qualche modo?

Fatemi sapere se qualcuno ha bisogno di più dettagli.

PS

Sto immaginando la soluzione è qualcosa di simile a quello che JAX non autograd, che è ciò che io sono più familiarità con. Quello che voglio dire qui è che JAX credo che si possa fare solo:

from jax import grad
df = grad(lambda x: f(theta, x))

e poi df potrebbe essere semplicemente una funzione che può essere chiamato in qualsiasi punto. Ma è PyTorch lo stesso? O c'è qualche conflitto interno .backward() che la causa di questo errore?