Difficulties in using jacobian of torch.autograd.functional

I am solving PDE, so I need the jacobian matrix of the residual with respect to variables.

The math is shown in the picture

I want the vector residual to be differentiated by [pgnew[1,:,:],swnew[0,:,:],pgnew[2,:,:],swnew[1,:,:],pgnew[3,:,:],swnew[2,:,:]]

Here is my code

import torch
from torch.autograd.functional import jacobian
def get_residual (pgnew, swnew):
    residual_w  = 5*(swnew-swold)+T_w*((pgnew[2:,:,:]-pgnew[1:-1,:,:])-(pc[2:,:,:]-pc[1:-1,:,:])) - T_w*((pgnew[1:-1,:,:]-pgnew[0:-2,:,:])-(pc[1:-1,:,:]-pc[0:-2,:,:]))
    residual_g  = 5*((1-swnew)-(1-swold))+T_g*(pgnew[2:,:,:]-pgnew[1:-1,:,:]) - T_g*(pgnew[1:-1,:,:]-pgnew[0:-2,:,:])
    residual    = torch.ravel(torch.column_stack((residual_w,residual_g)))
    return residual

if __name__ == '__main__':
    dt       = 0.01
    T_w      = 10 
    T_g      = 12
    swnew     = torch.zeros(3, 1, 1, requires_grad=True, dtype=torch.float64)   
    swold     = torch.ones(3, 1, 1, requires_grad=True, dtype=torch.float64)     
    pgnew     = 2*torch.ones(5, 1, 1, requires_grad=True, dtype=torch.float64)
    pc        = 3*torch.ones(5, 1, 1, requires_grad=True, dtype=torch.float64)
    

    unknown  = torch.ravel(torch.column_stack((pgnew[1:-1,:,:],swnew)))
    
    residual = get_residual(pgnew, swnew)
    print('Check Jacobian \n', jacobian(get_residual, unknown))

I am following this tutorial Jacobian matrix in PyTorch - GeeksforGeeks
However, it shows an error, namely,

get_residual() missing 1 required positional argument: ‘swnew’,

so I change it to

 print('Check Jacobian \n', jacobian(get_residual(pgnew, swnew),unknown)

Then it shows ‘Tensor’ object is not callable’

Thank you.

Hi Miraboreasu!

jacobian() is calling your function get_residual(). How jacobian() calls
get_residual() depends on the second argument you pass to jacobian().
You are calling jacobian() with the single tensor unknown as its second
argument, so jacobian() calls get_residual() with just a single tensor
argument. But your get_residual() function requires two arguments, hence
the error.

jacobian() requires its first argument to be a function (more precisely, a
“callable” – something that can be called like a function). But you are
calling jacobian() not with a function, but with the result of calling
get_residual (pgnew, swnew). Your get_residual() returns a Tensor
(which isn’t a callable), hence the error.

Best.

K. Frank

1 Like