prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN ISSUE ---
  QoL: better print for FrozenDict
  The best way I'm aware of to get an overview of model shape is via `jax.tree_map(jnp.shape, params)`. FrozenDicts have no concept of pretty printing the way dicts do, so large models are unwieldy to parse at a glance. 
  
  --- END ISSUE ---
  
  Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
  --- BEGIN FILES ---
  Path: `flax/core/frozen_dict.py`
  Content:
  ```
  1 # Copyright 2020 The Flax Authors.
  2 #
  3 # Licensed under the Apache License, Version 2.0 (the "License");
  4 # you may not use this file except in compliance with the License.
  5 # You may obtain a copy of the License at
  6 #
  7 #     http://www.apache.org/licenses/LICENSE-2.0
  8 #
  9 # Unless required by applicable law or agreed to in writing, software
  10 # distributed under the License is distributed on an "AS IS" BASIS,
  11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12 # See the License for the specific language governing permissions and
  13 # limitations under the License.
  14 
  15 """Frozen Dictionary."""
  16 
  17 from typing import Any, TypeVar, Mapping, Dict, Tuple
  18 
  19 from flax import serialization
  20 import jax
  21 
  22 
  23 K = TypeVar('K')
  24 V = TypeVar('V')
  25 
  26 
  27 @jax.tree_util.register_pytree_node_class
  28 class FrozenDict(Mapping[K, V]):
  29   """An immutable variant of the Python dict."""
  30   __slots__ = ('_dict', '_hash')
  31 
  32   def __init__(self, *args, **kwargs):
  33     # make sure the dict is as
  34     xs = dict(*args, **kwargs)
  35     self._dict = _prepare_freeze(xs)
  36 
  37     self._hash = None
  38 
  39   def __getitem__(self, key):
  40     v = self._dict[key]
  41     if isinstance(v, dict):
  42       return FrozenDict(v)
  43     return v
  44 
  45   def __setitem__(self, key, value):
  46     raise ValueError('FrozenDict is immutable.')
  47 
  48   def __contains__(self, key):
  49     return key in self._dict
  50 
  51   def __iter__(self):
  52     return iter(self._dict)
  53 
  54   def __len__(self):
  55     return len(self._dict)
  56 
  57   def __repr__(self):
  58     return 'FrozenDict(%r)' % self._dict
  59 
  60   def __hash__(self):
  61     if self._hash is None:
  62       h = 0
  63       for key, value in self.items():
  64         h ^= hash((key, value))
  65       self._hash = h
  66     return self._hash
  67 
  68   def copy(self, add_or_replace: Mapping[K, V]) -> 'FrozenDict[K, V]':
  69     """Create a new FrozenDict with additional or replaced entries."""
  70     return type(self)(self, **unfreeze(add_or_replace))
  71 
  72   def items(self):
  73     for key in self._dict:
  74       yield (key, self[key])
  75 
  76   def pop(self, key: K) -> Tuple['FrozenDict[K, V]', V]:
  77     """Create a new FrozenDict where one entry is removed.
  78 
  79     Example::
  80 
  81       state, params = variables.pop('params')
  82 
  83     Args:
  84       key: the key to remove from the dict
  85     Returns:
  86       A pair with the new FrozenDict and the removed value.
  87     """
  88     value = self[key]
  89     new_dict = dict(self._dict)
  90     new_dict.pop(key)
  91     new_self = type(self)(new_dict)
  92     return new_self, value
  93 
  94   def unfreeze(self) -> Dict[K, V]:
  95     return unfreeze(self)
  96 
  97   def tree_flatten(self):
  98     return (self._dict,), ()
  99 
  100   @classmethod
  101   def tree_unflatten(cls, _, data):
  102     return cls(*data)
  103 
  104 
  105 def _prepare_freeze(xs: Any) -> Any:
  106   """Deep copy unfrozen dicts to make the dictionary FrozenDict safe."""
  107   if isinstance(xs, FrozenDict):
  108     # we can safely ref share the internal state of a FrozenDict
  109     # because it is immutable.
  110     return xs._dict  # pylint: disable=protected-access
  111   if not isinstance(xs, dict):
  112     # return a leaf as is.
  113     return xs
  114   # recursively copy dictionary to avoid ref sharing
  115   return {key: _prepare_freeze(val) for key, val in xs.items()}
  116 
  117 
  118 def freeze(xs: Dict[K, V]) -> FrozenDict[K, V]:
  119   """Freeze a nested dict.
  120 
  121   Makes a nested `dict` immutable by transforming it into `FrozenDict`.
  122   """
  123   return FrozenDict(xs)
  124 
  125 
  126 def unfreeze(x: FrozenDict[K, V]) -> Dict[K, V]:
  127   """Unfreeze a FrozenDict.
  128 
  129   Makes a mutable copy of a `FrozenDict` mutable by transforming
  130   it into (nested) dict.
  131   """
  132   if not isinstance(x, (FrozenDict, dict)):
  133     return x
  134   ys = {}
  135   for key, value in x.items():
  136     ys[key] = unfreeze(value)
  137   return ys
  138 
  139 
  140 def _frozen_dict_state_dict(xs):
  141   return {key: serialization.to_state_dict(value) for key, value in xs.items()}
  142 
  143 
  144 def _restore_frozen_dict(xs, states):
  145   return FrozenDict(
  146       {key: serialization.from_state_dict(value, states[key])
  147        for key, value in xs.items()})
  148 
  149 
  150 serialization.register_serialization_state(
  151     FrozenDict,
  152     _frozen_dict_state_dict,
  153     _restore_frozen_dict)
  154
  ```
  --- END FILES ---
  
  Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.
  
  Here is an example:
  
  ```diff
  diff --git a/examples/server_async.py b/examples/server_async.py
  --- a/examples/server_async.py
  +++ b/examples/server_async.py
  @@ -313,4 +313,4 @@
   
   
   if __name__ == "__main__":
  -    asyncio.run(run_async_server("."), debug=True)
  +    asyncio.run(run_async_server(), debug=True)
  diff --git a/examples/server_sync.py b/examples/server_sync.py
  --- a/examples/server_sync.py
  +++ b/examples/server_sync.py
  @@ -313,5 +313,5 @@
   
   
   if __name__ == "__main__":
  -    server = run_sync_server(".")
  +    server = run_sync_server()
       server.shutdown()
  
  ```
  
golden_diff: diff --git a/flax/core/frozen_dict.py b/flax/core/frozen_dict.py
  --- a/flax/core/frozen_dict.py
  +++ b/flax/core/frozen_dict.py
  @@ -24,6 +24,14 @@
   V = TypeVar('V')
   
   
  +def _indent(x, num_spaces):
  +  indent_str = ' ' * num_spaces
  +  lines = x.split('\n')
  +  assert lines[-1] == ''
  +  # skip the final line because it's empty and should not be indented.
  +  return '\n'.join(indent_str + line for line in lines[:-1]) + '\n'
  +
  +
   @jax.tree_util.register_pytree_node_class
   class FrozenDict(Mapping[K, V]):
     """An immutable variant of the Python dict."""
  @@ -55,7 +63,21 @@
       return len(self._dict)
   
     def __repr__(self):
  -    return 'FrozenDict(%r)' % self._dict
  +    return self.pretty_repr()
  +
  +  def pretty_repr(self, num_spaces=4):
  +    """Returns an indented representation of the nested dictionary."""
  +    def pretty_dict(x):
  +      if not isinstance(x, dict):
  +        return repr(x)
  +      rep = ''
  +      for key, val in x.items():
  +        rep += f'{key}: {pretty_dict(val)},\n'
  +      if rep:
  +        return '{\n' + _indent(rep, num_spaces) + '}'
  +      else:
  +        return '{}'
  +    return f'FrozenDict({pretty_dict(self._dict)})'
   
     def __hash__(self):
       if self._hash is None:
  
verification_info: {"golden_diff": "diff --git a/flax/core/frozen_dict.py b/flax/core/frozen_dict.py\n--- a/flax/core/frozen_dict.py\n+++ b/flax/core/frozen_dict.py\n@@ -24,6 +24,14 @@\n V = TypeVar('V')\n \n \n+def _indent(x, num_spaces):\n+  indent_str = ' ' * num_spaces\n+  lines = x.split('\\n')\n+  assert lines[-1] == ''\n+  # skip the final line because it's empty and should not be indented.\n+  return '\\n'.join(indent_str + line for line in lines[:-1]) + '\\n'\n+\n+\n @jax.tree_util.register_pytree_node_class\n class FrozenDict(Mapping[K, V]):\n   \"\"\"An immutable variant of the Python dict.\"\"\"\n@@ -55,7 +63,21 @@\n     return len(self._dict)\n \n   def __repr__(self):\n-    return 'FrozenDict(%r)' % self._dict\n+    return self.pretty_repr()\n+\n+  def pretty_repr(self, num_spaces=4):\n+    \"\"\"Returns an indented representation of the nested dictionary.\"\"\"\n+    def pretty_dict(x):\n+      if not isinstance(x, dict):\n+        return repr(x)\n+      rep = ''\n+      for key, val in x.items():\n+        rep += f'{key}: {pretty_dict(val)},\\n'\n+      if rep:\n+        return '{\\n' + _indent(rep, num_spaces) + '}'\n+      else:\n+        return '{}'\n+    return f'FrozenDict({pretty_dict(self._dict)})'\n \n   def __hash__(self):\n     if self._hash is None:\n", "issue": "QoL: better print for FrozenDict\nThe best way I'm aware of to get an overview of model shape is via `jax.tree_map(jnp.shape, params)`. FrozenDicts have no concept of pretty printing the way dicts do, so large models are unwieldy to parse at a glance. \n", "before_files": [{"content": "# Copyright 2020 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Frozen Dictionary.\"\"\"\n\nfrom typing import Any, TypeVar, Mapping, Dict, Tuple\n\nfrom flax import serialization\nimport jax\n\n\nK = TypeVar('K')\nV = TypeVar('V')\n\n\n@jax.tree_util.register_pytree_node_class\nclass FrozenDict(Mapping[K, V]):\n  \"\"\"An immutable variant of the Python dict.\"\"\"\n  __slots__ = ('_dict', '_hash')\n\n  def __init__(self, *args, **kwargs):\n    # make sure the dict is as\n    xs = dict(*args, **kwargs)\n    self._dict = _prepare_freeze(xs)\n\n    self._hash = None\n\n  def __getitem__(self, key):\n    v = self._dict[key]\n    if isinstance(v, dict):\n      return FrozenDict(v)\n    return v\n\n  def __setitem__(self, key, value):\n    raise ValueError('FrozenDict is immutable.')\n\n  def __contains__(self, key):\n    return key in self._dict\n\n  def __iter__(self):\n    return iter(self._dict)\n\n  def __len__(self):\n    return len(self._dict)\n\n  def __repr__(self):\n    return 'FrozenDict(%r)' % self._dict\n\n  def __hash__(self):\n    if self._hash is None:\n      h = 0\n      for key, value in self.items():\n        h ^= hash((key, value))\n      self._hash = h\n    return self._hash\n\n  def copy(self, add_or_replace: Mapping[K, V]) -> 'FrozenDict[K, V]':\n    \"\"\"Create a new FrozenDict with additional or replaced entries.\"\"\"\n    return type(self)(self, **unfreeze(add_or_replace))\n\n  def items(self):\n    for key in self._dict:\n      yield (key, self[key])\n\n  def pop(self, key: K) -> Tuple['FrozenDict[K, V]', V]:\n    \"\"\"Create a new FrozenDict where one entry is removed.\n\n    Example::\n\n      state, params = variables.pop('params')\n\n    Args:\n      key: the key to remove from the dict\n    Returns:\n      A pair with the new FrozenDict and the removed value.\n    \"\"\"\n    value = self[key]\n    new_dict = dict(self._dict)\n    new_dict.pop(key)\n    new_self = type(self)(new_dict)\n    return new_self, value\n\n  def unfreeze(self) -> Dict[K, V]:\n    return unfreeze(self)\n\n  def tree_flatten(self):\n    return (self._dict,), ()\n\n  @classmethod\n  def tree_unflatten(cls, _, data):\n    return cls(*data)\n\n\ndef _prepare_freeze(xs: Any) -> Any:\n  \"\"\"Deep copy unfrozen dicts to make the dictionary FrozenDict safe.\"\"\"\n  if isinstance(xs, FrozenDict):\n    # we can safely ref share the internal state of a FrozenDict\n    # because it is immutable.\n    return xs._dict  # pylint: disable=protected-access\n  if not isinstance(xs, dict):\n    # return a leaf as is.\n    return xs\n  # recursively copy dictionary to avoid ref sharing\n  return {key: _prepare_freeze(val) for key, val in xs.items()}\n\n\ndef freeze(xs: Dict[K, V]) -> FrozenDict[K, V]:\n  \"\"\"Freeze a nested dict.\n\n  Makes a nested `dict` immutable by transforming it into `FrozenDict`.\n  \"\"\"\n  return FrozenDict(xs)\n\n\ndef unfreeze(x: FrozenDict[K, V]) -> Dict[K, V]:\n  \"\"\"Unfreeze a FrozenDict.\n\n  Makes a mutable copy of a `FrozenDict` mutable by transforming\n  it into (nested) dict.\n  \"\"\"\n  if not isinstance(x, (FrozenDict, dict)):\n    return x\n  ys = {}\n  for key, value in x.items():\n    ys[key] = unfreeze(value)\n  return ys\n\n\ndef _frozen_dict_state_dict(xs):\n  return {key: serialization.to_state_dict(value) for key, value in xs.items()}\n\n\ndef _restore_frozen_dict(xs, states):\n  return FrozenDict(\n      {key: serialization.from_state_dict(value, states[key])\n       for key, value in xs.items()})\n\n\nserialization.register_serialization_state(\n    FrozenDict,\n    _frozen_dict_state_dict,\n    _restore_frozen_dict)\n", "path": "flax/core/frozen_dict.py"}], "after_files": [{"content": "# Copyright 2020 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Frozen Dictionary.\"\"\"\n\nfrom typing import Any, TypeVar, Mapping, Dict, Tuple\n\nfrom flax import serialization\nimport jax\n\n\nK = TypeVar('K')\nV = TypeVar('V')\n\n\ndef _indent(x, num_spaces):\n  indent_str = ' ' * num_spaces\n  lines = x.split('\\n')\n  assert lines[-1] == ''\n  # skip the final line because it's empty and should not be indented.\n  return '\\n'.join(indent_str + line for line in lines[:-1]) + '\\n'\n\n\n@jax.tree_util.register_pytree_node_class\nclass FrozenDict(Mapping[K, V]):\n  \"\"\"An immutable variant of the Python dict.\"\"\"\n  __slots__ = ('_dict', '_hash')\n\n  def __init__(self, *args, **kwargs):\n    # make sure the dict is as\n    xs = dict(*args, **kwargs)\n    self._dict = _prepare_freeze(xs)\n\n    self._hash = None\n\n  def __getitem__(self, key):\n    v = self._dict[key]\n    if isinstance(v, dict):\n      return FrozenDict(v)\n    return v\n\n  def __setitem__(self, key, value):\n    raise ValueError('FrozenDict is immutable.')\n\n  def __contains__(self, key):\n    return key in self._dict\n\n  def __iter__(self):\n    return iter(self._dict)\n\n  def __len__(self):\n    return len(self._dict)\n\n  def __repr__(self):\n    return self.pretty_repr()\n\n  def pretty_repr(self, num_spaces=4):\n    \"\"\"Returns an indented representation of the nested dictionary.\"\"\"\n    def pretty_dict(x):\n      if not isinstance(x, dict):\n        return repr(x)\n      rep = ''\n      for key, val in x.items():\n        rep += f'{key}: {pretty_dict(val)},\\n'\n      if rep:\n        return '{\\n' + _indent(rep, num_spaces) + '}'\n      else:\n        return '{}'\n    return f'FrozenDict({pretty_dict(self._dict)})'\n\n  def __hash__(self):\n    if self._hash is None:\n      h = 0\n      for key, value in self.items():\n        h ^= hash((key, value))\n      self._hash = h\n    return self._hash\n\n  def copy(self, add_or_replace: Mapping[K, V]) -> 'FrozenDict[K, V]':\n    \"\"\"Create a new FrozenDict with additional or replaced entries.\"\"\"\n    return type(self)(self, **unfreeze(add_or_replace))\n\n  def items(self):\n    for key in self._dict:\n      yield (key, self[key])\n\n  def pop(self, key: K) -> Tuple['FrozenDict[K, V]', V]:\n    \"\"\"Create a new FrozenDict where one entry is removed.\n\n    Example::\n\n      state, params = variables.pop('params')\n\n    Args:\n      key: the key to remove from the dict\n    Returns:\n      A pair with the new FrozenDict and the removed value.\n    \"\"\"\n    value = self[key]\n    new_dict = dict(self._dict)\n    new_dict.pop(key)\n    new_self = type(self)(new_dict)\n    return new_self, value\n\n  def unfreeze(self) -> Dict[K, V]:\n    return unfreeze(self)\n\n  def tree_flatten(self):\n    return (self._dict,), ()\n\n  @classmethod\n  def tree_unflatten(cls, _, data):\n    return cls(*data)\n\n\ndef _prepare_freeze(xs: Any) -> Any:\n  \"\"\"Deep copy unfrozen dicts to make the dictionary FrozenDict safe.\"\"\"\n  if isinstance(xs, FrozenDict):\n    # we can safely ref share the internal state of a FrozenDict\n    # because it is immutable.\n    return xs._dict  # pylint: disable=protected-access\n  if not isinstance(xs, dict):\n    # return a leaf as is.\n    return xs\n  # recursively copy dictionary to avoid ref sharing\n  return {key: _prepare_freeze(val) for key, val in xs.items()}\n\n\ndef freeze(xs: Dict[K, V]) -> FrozenDict[K, V]:\n  \"\"\"Freeze a nested dict.\n\n  Makes a nested `dict` immutable by transforming it into `FrozenDict`.\n  \"\"\"\n  return FrozenDict(xs)\n\n\ndef unfreeze(x: FrozenDict[K, V]) -> Dict[K, V]:\n  \"\"\"Unfreeze a FrozenDict.\n\n  Makes a mutable copy of a `FrozenDict` mutable by transforming\n  it into (nested) dict.\n  \"\"\"\n  if not isinstance(x, (FrozenDict, dict)):\n    return x\n  ys = {}\n  for key, value in x.items():\n    ys[key] = unfreeze(value)\n  return ys\n\n\ndef _frozen_dict_state_dict(xs):\n  return {key: serialization.to_state_dict(value) for key, value in xs.items()}\n\n\ndef _restore_frozen_dict(xs, states):\n  return FrozenDict(\n      {key: serialization.from_state_dict(value, states[key])\n       for key, value in xs.items()})\n\n\nserialization.register_serialization_state(\n    FrozenDict,\n    _frozen_dict_state_dict,\n    _restore_frozen_dict)\n", "path": "flax/core/frozen_dict.py"}]}