Allowing users to create their own shapes

[TevenLeScao]

This feature is needed for users that want to define the get_shapes function themselves, for example in tensor-parallel environments where each process only sees a fraction of the parameters.