Fixes #1925 non-str dict keys not suppoted in module state

[muhrin]

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.

What does this PR do?

Fixes #1925

Checklist

• This PR fixes a minor issue (e.g.: typo or small bug) or improves the docs.
• This change is discussed in a Github issue/
  discussion .
• The documentation and docstrings adhere to the
  documentation guidelines.
• This change includes necessary high-coverage tests.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.