Instead of using torch.autograd, I’ve tried computing the Jacobian “manually”.
def manual_jacobian(x, model):
    # x -> fc1 -> s1
    # s1 -> relu -> a1
    # a1 -> fc2 -> s2 (output)

    register_hooks(model)
    output = model(x)

    n = output.size(0)
    fc1 = model.fc1
    fc2 = model.fc2

    # Jacobian s2 (output) -> a1
    Ja1 = torch.stack([fc2.weight] * n)

    # ReLU derivative a1 -> s1
    d = (fc2.data_input > 0).type(fc2.data_input.dtype)

    # Jacobian s2 (output) -> s1
    Js1 = contract('bij,bj->bij', Ja1, d)

    # Jacobian s2 (output) -> x
    Jx = contract('bij,jk->bik', Js1, fc1.weight)

    return Jx

here is the example code 13.
By limiting the type of layer (Linear, ReLU, and Conv2d for the example code above), we can know how to compute the Jacobian efficiently without using torch.autograd nor the repeat trick.
For example, we can get the Jacobian for ConvNd by using ConvTransposeNd 4
Benchmark on 1 GPU: my approach (manual mode) vs reverse mode for a Linear NN
bs: 100
n_inputs: 10
n_outputs: 100
loop: 100
-------------
reverse mode: 0.459s
manual mode: 0.195s
(Jx_rev - Jx_man).max(): 0.0

and it’s much faster when it comes to a large number of outputs
bs: 100
n_inputs: 10
n_outputs: 1000
loop: 100
-------------
reverse mode: 18.032s
manual mode: 1.884s
(Jx_rev - Jx_man).max(): 0.0

For a CNN
bs: 1
input shape: torch.Size([3, 32, 32])
n_outputs: 100
loop: 1
-------------
reverse mode: 0.416s
manual mode: 0.232s
(Jx_rev - Jx_man).max(): 3.725290298461914e-09

To summarize the pros & cons of  manual mode:
pros

it can compute Jacobian in a single backpropagation for an arbitrary number of outputs.
compare to the repeat trick, it doesn’t require any copy of the inputs for forwarding.
it can evaluate the Jacobian layer by layer, so it’s supposed to work even with a large model, e.g., ResNet-50 for ImageNet classification.

cons

we need to figure out the way to calculate the Jacobian for each module we need for constructing a network.
we need to detect the type and the order of the modules in the network.
it would be more complicated for a network with skip connections.

I am trying to overcome these shortcomings of manual mode, and I hope this approach would help.