Add test for missing functionality of shape inference when applying a function.

This feature is necessary to be able to easily apply functions in graph mode
for things like unknown batch size or variable image size/feature maps.

Without it a user of a function with an input_signature in graph mode has to
manually compute the shape and set it so the rest of the model building can
take that into account.

PiperOrigin-RevId: 234133298