Nice, thank you!

This is maybe at the cost of performance because I feel like doing and undoing things that could be taken care of in those two functions.

I guess the performance impact won't be too bad, but we could measure it. It's nice that the change is so local.

With this change, Theta is applied in pixel space. This means the translation must be given in pixels, not as a fraction of the input size. As I argued before, this might not be ideal, as it will require a very different value range for translations compared to rotation and shearing. Also it might make it more difficult to process images of variable input size; there's no easy way to specify cropping the top left corner, for example.
So another solution might be to figure out the aspect ratio and scale either x_s or y_s (not both), depending on whether the width or height is smaller.
In any case, the operation space of Theta should be determined by an option given to the Layer constructor. Existing networks should just work as before. So we could have three alternatives for a new grid_space or grid_domain option: pixels (your new implementation), unit_square (the old implementation, staying the default), and unit_rect (a rectangle with the longer side bounded in [0,1] and the shorter side being shorter). I'm not very sure about the naming, but does this seem sensible to you?

Yes I agree with your design principles, I'll implement them when I have time