Encode dicts of params into text.

PiperOrigin-RevId: 491943111

Author: lingvo-bot

lingvo/core/BUILD

@@ -833,6 +833,7 @@ pytype_library(
         ":symbolic",
         # Implicit python proto dependency.
         "//lingvo:compat",
+        # Implicit numpy dependency.
         # Implicit typing_extensions dependency.
     ],
 )

lingvo/core/hyperparams.py

@@ -104,7 +104,8 @@ class _SortedDict(dict):
 
   def __repr__(self):
     return '{' + ', '.join(
-        '%r: %r' % item for item in sorted(self.items())) + '}'
+        '%r: %s' % (k, _QuoteString(v) if isinstance(v, str) else repr(v))
+        for k, v in sorted(self.items())) + '}'
 
 
 class _Param:
@@ -390,17 +391,18 @@ def _GetNested(self, name: str) -> Tuple[ParamsT, str]:
     for i, part in enumerate(parts[:-1]):
       # Get the value (nested Params object) associated with name 'part'.
       try:
-        if is_list := re.match(r'^(.+)\[(.+)\]$', part):
-          part = is_list.group(1)
-          list_index = int(is_list.group(2))
+        if is_list_or_dict := re.match(r'^(.+)\[(.+)\]$', part):
+          part = is_list_or_dict.group(1)
+          list_index = ast.literal_eval(is_list_or_dict.group(2))
         # pylint: disable=protected-access
         curr = curr._params[part].Get()
-        if is_list:
+        if is_list_or_dict:
           curr = curr[list_index]
       except KeyError:
         raise AttributeError('.'.join(parts[:i + 1]))
       assert isinstance(curr, Params), ('Cannot introspect %s for %s' %
                                         (type(curr), '.'.join(parts[:i + 1])))
+
     return curr, parts[-1]
 
   def Set(self: ParamsT, **kwargs: Any) -> ParamsT:
@@ -697,7 +699,7 @@ def _Visit(key: str, val: Any):
       elif isinstance(val, dict):
         if enter_fn(key, val):
           for k, v in val.items():
-            _Visit(_SubKey(key, k), v)
+            _Visit(f'{key}[\'{k}\']', v)
           exit_fn(key, val)
         else:
           visit_fn(key, val)
@@ -803,6 +805,9 @@ def _Enter(key: str, p: Any) -> bool:
           isinstance(x, tuple) and len(x) == 2 and isinstance(x[0], str) and
           isinstance(x[1], Params) for x in p)):
         return True
+      elif (isinstance(p, dict) and p and all(
+          isinstance(k, str) and isinstance(v, Params) for k, v in p.items())):
+        return True
       return False
 
     kv = {}

lingvo/core/hyperparams_test.py

@@ -424,8 +424,8 @@ def Visit(key, value):
     outer.Visit(Visit)
 
     self.assertListEqual(visit_keys, [
-        'inner.a[0][0]', 'inner.a[0][1][0]', 'inner.a[0][1][1]',
-        'inner.a[0][1][2][x]'
+        'inner[\'a\'][0][0]', 'inner[\'a\'][0][1][0]', 'inner[\'a\'][0][1][1]',
+        'inner[\'a\'][0][1][2][x]'
     ])
 
   def testToText(self):
@@ -446,6 +446,10 @@ def testToText(self):
     outer.Define('plain_dict', {'a': 10}, '')
     outer.Define('complex_dict', {'a': 10, 'b': inner}, '')
     outer.Define('complex_dict_escape', {'a': 'abc"\'\ndef'}, '')
+    outer.Define('complex_dict_with_params', {
+        'a': inner,
+        'b': inner.Copy()
+    }, '')
     outer.Define('some_class', complex(0, 1), '')
     outer.Define('optional_bool', None, '')
     outer.Define('enum', TestEnum.B, '')
@@ -460,7 +464,12 @@ def testToText(self):
         '\n' + outer.ToText(), r"""
 class : type/__main__/TestClass1
 complex_dict : {'a': 10, 'b': {'bar': 2.71, 'baz': 'hello'}}
-complex_dict_escape : {'a': 'abc"\'\ndef'}
+complex_dict_escape : {'a': 'abc"\'
+def'}
+complex_dict_with_params['a'].bar : 2.71
+complex_dict_with_params['a'].baz : 'hello'
+complex_dict_with_params['b'].bar : 2.71
+complex_dict_with_params['b'].baz : 'hello'
 dataclass : {'a': [42], 'b': 'float32'}
 dtype : float32
 dtype2 : int32
@@ -488,6 +497,10 @@ class : type/__main__/TestClass1
         dtype2 : float32
         inner.baz : 'world'
         # foo : 123
+        complex_dict_with_params['a'].bar : 2.71
+        complex_dict_with_params['a'].baz : 'world'
+        complex_dict_with_params['b'].bar : 2.71
+        complex_dict_with_params['b'].baz : 'hello'
         optional_bool : true
         list_of_params[0].bar : 2.72
         seqlen : [1, 2.0, '3', [4]]
@@ -507,7 +520,12 @@ class : type/__main__/TestClass2
         '\n' + outer.ToText(), r"""
 class : type/__main__/TestClass2
 complex_dict : {'a': 10, 'b': {'bar': 2.71, 'baz': 'world'}}
-complex_dict_escape : {'a': 'abc"\'\ndef'}
+complex_dict_escape : {'a': 'abc"\'
+def'}
+complex_dict_with_params['a'].bar : 2.71
+complex_dict_with_params['a'].baz : 'world'
+complex_dict_with_params['b'].bar : 2.71
+complex_dict_with_params['b'].baz : 'hello'
 dataclass : {'a': 27, 'b': 'int32'}
 dtype : float32
 dtype2 : float32