There's no reason why not two custom vmappable types cannot share the same spec_type. However, spec_types was a set, which can cause bugs / exceptions.
Suppose that I register two vmappable data_types sharing the same spec_type, and then unregister one of the two. Then, the spec_type is no longer in the set to support the second data_type. Also, an exception will be raised if I try to unregister the two vmappable types (the second call to spec_types.remove).
When unregistering a data type, instead of removing its spec_type from the set, we regenerate the set from the remaining vmappable types.
PiperOrigin-RevId: 737280270
466ef6a132