fold/unfold could be a start if I was trying to pass patches through the same linear layer, however patches have to pass through different weights, unlike cnn where filter weights are shared. Do  you have any further ideas?
Thank you for your reply,