flax

287eb7b6 - This CL relaxes the input type hints to resolve an invariance issue with Python type checkers (Pyright/MyPy).

This CL relaxes the input type hints to resolve an invariance issue with Python type checkers (Pyright/MyPy). Rationale: The current type definition strictly requires dict[str, int]. Because standard dictionaries are invariant, this definition rejects dictionaries keyed by StrEnum, despite StrEnum being a subclass of str. This forces downstream users to cast their data types or lose type safety. The example below illustrates the false positive error generated by the current signature: ```Python import enum class MyEnum(enum.StrEnum): # Minimal replacement of e.g. research.shared.configs.base.StrEnumEnum a = enum.auto() b = enum.auto() def func_we_control() -> dict[MyEnum, int]: # It's important to keep the return as MyEnum, so that other code knows the keys of this. return {MyEnum.a: 1} def func_we_dont_control(x: dict[str, int]) -> None: del x pass func_we_dont_control(func_we_control()) ``` Pyright Error: Type parameter "_KT@dict" is invariant, but "MyEnum" is not the same as "str" PiperOrigin-RevId: 855783481