[Mosaic GPU] Add an `AnyOf` constraint to layout inference.
This constraint allows specifying that an expression must be equal to one of a given set of constant values. Layout inference is updated to extract potential assignments from `AnyOf` constraints.
PiperOrigin-RevId: 895779141
