Oh wait, I see what you do in the proof now. You lift it first into a larger space for injection.

How about this then:

Consider the product space (no need to worry about finite dimensions).

where injective.

where surjective.

Now this seems to be general enough to work for all functions and not just linear ones. Did I make a mistake?