How to save memory while having a big Jacobian matirx


I am doing some computation using pytorch, there is a Jacobian matrix, J of a vector of equations, r, and I am using jacfwd to calculate it, then there will be a linear solving

        r      = get_residual(unknown)
        J      = jacfwd(get_residual)(unknown)
        update       = torch.linalg.solve(J, r)

I actually don’t need J to be explicit, how to use like COO or CSR matrix to save memory and speedup this process?


I don’t need the full matrix from jacfwd, so is there way to save it to a tmp and directlly make it COO or CSR?