flax

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

#5187

Open