Actually pytorch might not work since it’s using reverse mode auto differentiation (back propagation). Backprop is super inefficient when output number is much larger than input number, which is likely to be the case in calculating Jacobian (num_samples v.s. num_parameters).
The easiest option for now is JAX if you don’t want to write everything by hand; it only takes me 30 lines to define a MLP class.