Add support for use_static_shapes.

If false, the output is always padded to max_total_size
If true, the output is padded to the minimum of max_total_size and
max_size_per_class * num_classes
Updated unit tests to take this into account