Fixes #1925
In PR #2007, the issue of having a dict with non-string keys as linen
module parameters was fixed but the related issue of having non-string
keys as part of the state dict was not. Specifically, these cases did
not work:
a = {1: 2}
class Foo(nn.Module):
def setup(self):
self.a = FrozenDict(a) # int as key.
class Foo(nn.Module):
a: dict[int, int] # int as key.
this is because serialization for non-int keys was fixes for ``dict``
but not for ``FrozenDict``. This uses the code from ``dict`` to add the
same behaviour to ``FrozenDict``.
