We know that we can implement our own custom autograd functions by subclassing torch.autograd.Function and implementing the forward and backward passes which operate on Tensors. As shown in https://pytorch.org/docs/stable/autograd.html#function 3
class Exp(Function):

    @staticmethod
    def forward(ctx, i):
        result = i.exp()
        ctx.save_for_backward(result)
        return result

    @staticmethod
    def backward(ctx, grad_output):
        result, = ctx.saved_tensors
        return grad_output * result

The question is how can I call the backward function of operations in torch.nn.functional? Is this possible? And what is the simplest way?
For example, I want to call the backward function of torch.nn.functional.softmax and get the Jacobian matrix. I have tried torch.autograd.functional.jacobian, but it is experimental now and gets a wrong result when I apply it to softmax and a 3D tensor. Is there an existing backward function that I can directly call?