Batch-wise Matrix-Free LSQR

I am trying to solve a least square problem for a linear system Ax - b, where x is a 1D vector , b is also a 1D vector that contains two sub-vectors b1 and b2, and A contains two operators: the identity map I and the differentiation operator D. This can be written as

A =[I,  D]), b =[b1, b2])


x = b1, D(x) = b2

The operator D is written as a function that takes the 1st order derivative of the signal x. The scipy function


with first argument a function can solve the above problem matrix-free by the linear operator


I am wondering whether there are similar features for

Or is there any alternative way to achieve this using pytorch?