Implement ``custom_model_2d``

Since we have ``custom_model_1d``, is there any reason we can't have ``custom_model_2d``?
I like the way that you use `custom_model` as a function in the first example and as a decorator here. Out would be even better if there was a sentence in the text that both uses are possible.