You can calculate (subgraph) gradients during the forward pass. When doing this inside autograd.Function, you can plug precalculated gradients into the outer graph. Not sure if this will help with your case, but here is an illustrative snippet:
class LpdfGenAsymLaplace(torch.autograd.Function):
	@staticmethod
	def forward(ctx, x_in, p0, g_in):
		if any(ctx.needs_input_grad):
			with torch.enable_grad():
				x = x_in.detach().requires_grad_()
				g = g_in.detach().requires_grad_()
				lpdf = _pdf(x,p0,g)
				dLpdf_dX, dLpdf_dG = torch.autograd.grad(lpdf, (x,g), torch.ones_like(lpdf), retain_graph=True, create_graph=True)
				ctx.save_for_backward(dLpdf_dX, dLpdf_dG)
				ctx.mark_non_differentiable(p0)
				return lpdf
		else:
			return _pdf(x_in,p0,g_in)
		
	@staticmethod
	def backward(ctx, gr_in):
		dLpdf_dX, dLpdf_dG = ctx.saved_tensors
		return dLpdf_dX * gr_in, None, dLpdf_dG * gr_in