pyglove 0.4.5.dev202411132359__py3-none-any.whl → 0.4.5.dev202501250807__py3-none-any.whl

pyglove/core/symbolic/dict_test.py

@@ -11,16 +11,14 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Tests for pyglove.Dict."""
-
 import copy
 import inspect
 import io
 import pickle
 import unittest
 
-from pyglove.core import object_utils
 from pyglove.core import typing as pg_typing
+from pyglove.core import utils
 from pyglove.core.symbolic import base
 from pyglove.core.symbolic import flags
 from pyglove.core.symbolic import inferred
@@ -31,7 +29,7 @@ from pyglove.core.symbolic.pure_symbolic import NonDeterministic
 from pyglove.core.symbolic.pure_symbolic import PureSymbolic
 
 
-MISSING_VALUE = object_utils.MISSING_VALUE
+MISSING_VALUE = utils.MISSING_VALUE
 
 
 class DictTest(unittest.TestCase):
@@ -44,14 +42,14 @@ class DictTest(unittest.TestCase):
     self.assertEqual(len(sd), 0)
 
     # Schemaless dict created from a regular dict.
-    sd = Dict({'a': 1})
+    sd = Dict({'a': 1, 1: 1})
     self.assertIsNone(sd.value_spec)
-    self.assertEqual(sd, dict(a=1))
+    self.assertEqual(sd, {'a': 1, 1: 1})
 
     # Schemaless dict created from key value pairs.
-    sd = Dict((('a', 1),))
+    sd = Dict((('a', 1), (1, 1)))
     self.assertIsNone(sd.value_spec)
-    self.assertEqual(sd, dict(a=1))
+    self.assertEqual(sd, {'a': 1, 1: 1})
 
     # Schemaless dict created from keyword args.
     sd = Dict(a=1)
@@ -59,9 +57,9 @@ class DictTest(unittest.TestCase):
     self.assertEqual(sd, dict(a=1))
 
     # Schemaless dict created from both a regular dict and keyword args.
-    sd = Dict({'a': 1}, a=2)
+    sd = Dict({'a': 1, 1: 1}, a=2)
     self.assertIsNone(sd.value_spec)
-    self.assertEqual(sd, dict(a=2))
+    self.assertEqual(sd, {'a': 2, 1: 1})
 
     # Schematized dict.
     vs = pg_typing.Dict([('a', pg_typing.Int())])
@@ -305,8 +303,8 @@ class DictTest(unittest.TestCase):
       with self.assertRaisesRegex(
           base.WritePermissionError, 'Cannot modify field of a sealed Dict.'):
         sd['b'] = 1
-    with self.assertRaisesRegex(KeyError, 'Key must be string type'):
-      sd[0] = 1
+    with self.assertRaisesRegex(KeyError, 'Key must be string or int type'):
+      sd[0.5] = 1
 
     # Set item in a schematized dict.
     sd = Dict(value_spec=pg_typing.Dict([('a', pg_typing.Int(default=0))]))
@@ -786,35 +784,45 @@ class DictTest(unittest.TestCase):
         sd.missing_values(flatten=False), {'a': MISSING_VALUE})
 
     # A non-schematized dict has a schematized child.
-    sd = Dict(x=sd)
+    sd = Dict({'x': sd, 1: sd})
     self.assertIsNone(sd.value_spec)
-    self.assertEqual(sd.missing_values(), {'x.a': MISSING_VALUE})
+    self.assertEqual(
+        sd.missing_values(), {'x.a': MISSING_VALUE, '[1].a': MISSING_VALUE}
+    )
 
   def test_sym_has(self):
-    sd = Dict(x=1, y=Dict(z=2))
+    sd = Dict({'x': 1, 1: dict(a=3), 'y': Dict({'z': 2, 2: 3})})
     self.assertTrue(sd.sym_has('x'))
+    self.assertTrue(sd.sym_has(1))
     self.assertTrue(sd.sym_has('y.z'))
-    self.assertTrue(sd.sym_has(object_utils.KeyPath.parse('y.z')))
+    self.assertTrue(sd.sym_has('[1].a'))
+    self.assertTrue(sd.sym_has('y[2]'))
+    self.assertTrue(sd.sym_has(utils.KeyPath.parse('y.z')))
     self.assertFalse(sd.sym_has('x.z'))
 
   def test_sym_get(self):
-    sd = Dict(x=1, y=Dict(z=2))
+    sd = Dict({'x': 1, 1: dict(a=3), 'y': Dict({'z': 2, 2: 3})})
     self.assertEqual(sd.sym_get('x'), 1)
+    self.assertEqual(sd.sym_get(1), dict(a=3))
     self.assertEqual(sd.sym_get('y.z'), 2)
+    self.assertEqual(sd.sym_get('[1].a'), 3)
+    self.assertEqual(sd.sym_get('y[2]'), 3)
     self.assertIsNone(sd.sym_get('x.z', None))
     with self.assertRaisesRegex(
         KeyError, 'Cannot query sub-key \'z\' of object.'):
       sd.sym_get('x.z')
 
   def test_sym_hasattr(self):
-    sd = Dict(x=1, y=Dict(z=2))
+    sd = Dict({'x': 1, 1: dict(a=3), 'y': Dict({'z': 2, 2: 3})})
     self.assertTrue(sd.sym_hasattr('x'))
+    self.assertTrue(sd.sym_hasattr(1))
     self.assertFalse(sd.sym_hasattr('y.z'))
     self.assertFalse(sd.sym_hasattr('a'))
 
   def test_sym_getattr(self):
-    sd = Dict(x=1, y=Dict(z=2))
+    sd = Dict({'x': 1, 1: dict(a=3), 'y': Dict({'z': 2, 2: 3})})
     self.assertEqual(sd.sym_getattr('x'), 1)
+    self.assertEqual(sd.sym_getattr(1), dict(a=3))
     self.assertIsNone(sd.sym_getattr('a', None))
     with self.assertRaisesRegex(
         AttributeError,
@@ -832,8 +840,9 @@ class DictTest(unittest.TestCase):
     with self.assertRaisesRegex(AttributeError, 'z'):
       _ = sd.sym_inferred('z')
 
-    sd = Dict(y=1, x=Dict(x=Dict(y=inferred.ValueFromParentChain())))
+    sd = Dict(y=1, x={'x': Dict(y=inferred.ValueFromParentChain()), 1: 2})
     self.assertEqual(sd.x.x.y, 1)
+    self.assertEqual(sd.x[1], 2)
 
   def test_sym_field(self):
     sd = Dict(x=1, y=Dict(z=2))
@@ -859,9 +868,9 @@ class DictTest(unittest.TestCase):
     self.assertIs(sd.sym_attr_field('x'), spec.schema.get_field('x'))
 
   def test_sym_keys(self):
-    sd = Dict(x=1, y=2)
+    sd = Dict({'x': 1, 'y': 2, 1: 3})
     self.assertEqual(next(sd.sym_keys()), 'x')
-    self.assertEqual(list(sd.sym_keys()), ['x', 'y'])
+    self.assertEqual(list(sd.sym_keys()), ['x', 'y', 1])
 
     sd = Dict(x=1, z=3, y=2, value_spec=pg_typing.Dict([
         (pg_typing.StrKey(), pg_typing.Int())
@@ -874,9 +883,9 @@ class DictTest(unittest.TestCase):
     self.assertEqual(list(sd.sym_keys()), ['x', 'y'])
 
   def test_sym_values(self):
-    sd = Dict(x=1, y=2)
+    sd = Dict({'x': 1, 'y': 2, 1: 3})
     self.assertEqual(next(sd.sym_values()), 1)
-    self.assertEqual(list(sd.sym_values()), [1, 2])
+    self.assertEqual(list(sd.sym_values()), [1, 2, 3])
 
     sd = Dict(x=1, z=3, y=2, value_spec=pg_typing.Dict([
         (pg_typing.StrKey(), pg_typing.Int())
@@ -891,9 +900,9 @@ class DictTest(unittest.TestCase):
     )
 
   def test_sym_items(self):
-    sd = Dict(x=1, y=2)
+    sd = Dict({'x': 1, 'y': 2, 1: 3})
     self.assertEqual(next(sd.sym_items()), ('x', 1))
-    self.assertEqual(list(sd.sym_items()), [('x', 1), ('y', 2)])
+    self.assertEqual(list(sd.sym_items()), [('x', 1), ('y', 2), (1, 3)])
 
     sd = Dict(x=1, z=3, y=2, value_spec=pg_typing.Dict([
         (pg_typing.StrKey(), pg_typing.Int())
@@ -917,9 +926,9 @@ class DictTest(unittest.TestCase):
 
   def test_sym_rebind(self):
     # Refer to RebindTest for more detailed tests.
-    sd = Dict(x=1, y=2)
-    sd.sym_rebind(x=2)
-    self.assertEqual(sd, dict(x=2, y=2))
+    sd = Dict({'x': 1, 'y': 2, 1: 3})
+    sd.sym_rebind({1: 4}, x=2)
+    self.assertEqual(sd, {'x': 2, 'y': 2, 1: 4})
 
   def test_sym_clone(self):
     class A:
@@ -997,6 +1006,12 @@ class DictTest(unittest.TestCase):
       y: int = 1
       use_symbolic_comparison = True
 
+    sd = Dict({'x': A(2), 1: B(2)})
+    self.assertEqual(sd.sym_nondefault(), {'x.x': 2, '[1].y': 2})
+
+    sd = Dict({'x': A(1), 1: B(1)})
+    self.assertEqual(sd.sym_nondefault(), {'x.x': 1})
+
     sd = Dict(x=1, y=dict(a1=A(1)), value_spec=pg_typing.Dict([
         ('x', pg_typing.Int(default=0)),
         ('y', pg_typing.Dict([
@@ -1080,7 +1095,7 @@ class DictTest(unittest.TestCase):
     self.assertTrue(Dict().sym_eq(Dict()))
     self.assertTrue(base.eq(Dict(), Dict()))
 
-    self.assertEqual(Dict(a=1), Dict(a=1))
+    self.assertEqual(Dict({'a': 1, 1: 2}), Dict({'a': 1, 1: 2}))
     self.assertTrue(Dict(a=1).sym_eq(Dict(a=1)))
     self.assertTrue(base.eq(Dict(a=1), Dict(a=1)))
     self.assertTrue(
@@ -1130,6 +1145,8 @@ class DictTest(unittest.TestCase):
     self.assertTrue(base.ne(Dict(), 1))
     self.assertNotEqual(Dict(), Dict(a=1))
     self.assertTrue(base.ne(Dict(), Dict(a=1)))
+    self.assertNotEqual(Dict({1: 1}), Dict({'1': 1}))
+    self.assertTrue(base.ne(Dict({1: 1}), Dict({'1': 1})))
     self.assertNotEqual(Dict(a=0), Dict(a=1))
     self.assertTrue(base.ne(Dict(a=0), Dict(a=1)))
 
@@ -1192,6 +1209,7 @@ class DictTest(unittest.TestCase):
     self.assertEqual(hash(Dict(a=1)), hash(Dict(a=1)))
     self.assertEqual(hash(Dict(a=dict(x=1))), hash(Dict(a=dict(x=1))))
     self.assertNotEqual(hash(Dict()), hash(Dict(a=1)))
+    self.assertNotEqual(hash(Dict({'1': 1})), hash(Dict({1: 1})))
     self.assertNotEqual(hash(Dict(a=1)), hash(Dict(a=2)))
 
     class A:
@@ -1247,7 +1265,7 @@ class DictTest(unittest.TestCase):
     self.assertEqual(sd.x.a.sym_path, 'x.a')
     self.assertEqual(sd.y[0].b.sym_path, 'y[0].b')
 
-    sd.sym_setpath(object_utils.KeyPath('a'))
+    sd.sym_setpath(utils.KeyPath('a'))
     self.assertEqual(sd.sym_path, 'a')
     self.assertEqual(sd.x.sym_path, 'a.x')
     self.assertEqual(sd.x.a.sym_path, 'a.x.a')
@@ -1704,96 +1722,112 @@ class RebindTest(unittest.TestCase):
         'd': 'foo',  # Unchanged.
         'e': 'bar'
     })
-    self.assertEqual(updates, [
-        {  # Notification to `sd.c[0]`.
-            'p': base.FieldUpdate(
-                object_utils.KeyPath.parse('c[0].p'),
-                target=sd.c[0],
-                field=None,
-                old_value=1,
-                new_value=MISSING_VALUE),
-            'q': base.FieldUpdate(
-                object_utils.KeyPath.parse('c[0].q'),
-                target=sd.c[0],
-                field=None,
-                old_value=MISSING_VALUE,
-                new_value=2),
-        },
-        {  # Notification to `sd.c`.
-            '[0].p': base.FieldUpdate(
-                object_utils.KeyPath.parse('c[0].p'),
-                target=sd.c[0],
-                field=None,
-                old_value=1,
-                new_value=MISSING_VALUE),
-            '[0].q': base.FieldUpdate(
-                object_utils.KeyPath.parse('c[0].q'),
-                target=sd.c[0],
-                field=None,
-                old_value=MISSING_VALUE,
-                new_value=2),
-        },
-        {  # Notification to `sd.b.y`.
-            'z': base.FieldUpdate(
-                object_utils.KeyPath.parse('b.y.z'),
-                target=sd.b.y,
-                field=None,
-                old_value=MISSING_VALUE,
-                new_value=1),
-        },
-        {  # Notification to `sd.b`.
-            'x': base.FieldUpdate(
-                object_utils.KeyPath.parse('b.x'),
-                target=sd.b,
-                field=None,
-                old_value=1,
-                new_value=2),
-            'y.z': base.FieldUpdate(
-                object_utils.KeyPath.parse('b.y.z'),
-                target=sd.b.y,
-                field=None,
-                old_value=MISSING_VALUE,
-                new_value=1),
-        },
-        {  # Notification to `sd`.
-            'a': base.FieldUpdate(
-                object_utils.KeyPath.parse('a'),
-                target=sd,
-                field=None,
-                old_value=1,
-                new_value=2),
-            'b.x': base.FieldUpdate(
-                object_utils.KeyPath.parse('b.x'),
-                target=sd.b,
-                field=None,
-                old_value=1,
-                new_value=2),
-            'b.y.z': base.FieldUpdate(
-                object_utils.KeyPath.parse('b.y.z'),
-                target=sd.b.y,
-                field=None,
-                old_value=MISSING_VALUE,
-                new_value=1),
-            'c[0].p': base.FieldUpdate(
-                object_utils.KeyPath.parse('c[0].p'),
-                target=sd.c[0],
-                field=None,
-                old_value=1,
-                new_value=MISSING_VALUE),
-            'c[0].q': base.FieldUpdate(
-                object_utils.KeyPath.parse('c[0].q'),
-                target=sd.c[0],
-                field=None,
-                old_value=MISSING_VALUE,
-                new_value=2),
-            'e': base.FieldUpdate(
-                object_utils.KeyPath.parse('e'),
-                target=sd,
-                field=None,
-                old_value=MISSING_VALUE,
-                new_value='bar')
-        }
-    ])
+    self.assertEqual(
+        updates,
+        [
+            {  # Notification to `sd.c[0]`.
+                'p': base.FieldUpdate(
+                    utils.KeyPath.parse('c[0].p'),
+                    target=sd.c[0],
+                    field=None,
+                    old_value=1,
+                    new_value=MISSING_VALUE,
+                ),
+                'q': base.FieldUpdate(
+                    utils.KeyPath.parse('c[0].q'),
+                    target=sd.c[0],
+                    field=None,
+                    old_value=MISSING_VALUE,
+                    new_value=2,
+                ),
+            },
+            {  # Notification to `sd.c`.
+                '[0].p': base.FieldUpdate(
+                    utils.KeyPath.parse('c[0].p'),
+                    target=sd.c[0],
+                    field=None,
+                    old_value=1,
+                    new_value=MISSING_VALUE,
+                ),
+                '[0].q': base.FieldUpdate(
+                    utils.KeyPath.parse('c[0].q'),
+                    target=sd.c[0],
+                    field=None,
+                    old_value=MISSING_VALUE,
+                    new_value=2,
+                ),
+            },
+            {  # Notification to `sd.b.y`.
+                'z': base.FieldUpdate(
+                    utils.KeyPath.parse('b.y.z'),
+                    target=sd.b.y,
+                    field=None,
+                    old_value=MISSING_VALUE,
+                    new_value=1,
+                ),
+            },
+            {  # Notification to `sd.b`.
+                'x': base.FieldUpdate(
+                    utils.KeyPath.parse('b.x'),
+                    target=sd.b,
+                    field=None,
+                    old_value=1,
+                    new_value=2,
+                ),
+                'y.z': base.FieldUpdate(
+                    utils.KeyPath.parse('b.y.z'),
+                    target=sd.b.y,
+                    field=None,
+                    old_value=MISSING_VALUE,
+                    new_value=1,
+                ),
+            },
+            {  # Notification to `sd`.
+                'a': base.FieldUpdate(
+                    utils.KeyPath.parse('a'),
+                    target=sd,
+                    field=None,
+                    old_value=1,
+                    new_value=2,
+                ),
+                'b.x': base.FieldUpdate(
+                    utils.KeyPath.parse('b.x'),
+                    target=sd.b,
+                    field=None,
+                    old_value=1,
+                    new_value=2,
+                ),
+                'b.y.z': base.FieldUpdate(
+                    utils.KeyPath.parse('b.y.z'),
+                    target=sd.b.y,
+                    field=None,
+                    old_value=MISSING_VALUE,
+                    new_value=1,
+                ),
+                'c[0].p': base.FieldUpdate(
+                    utils.KeyPath.parse('c[0].p'),
+                    target=sd.c[0],
+                    field=None,
+                    old_value=1,
+                    new_value=MISSING_VALUE,
+                ),
+                'c[0].q': base.FieldUpdate(
+                    utils.KeyPath.parse('c[0].q'),
+                    target=sd.c[0],
+                    field=None,
+                    old_value=MISSING_VALUE,
+                    new_value=2,
+                ),
+                'e': base.FieldUpdate(
+                    utils.KeyPath.parse('e'),
+                    target=sd,
+                    field=None,
+                    old_value=MISSING_VALUE,
+                    new_value='bar',
+                ),
+            },
+        ],
+    )
 
   def test_rebind_with_fn(self):
     sd = Dict(a=1, b=dict(x=2, y='foo', z=[0, 1, 2]))
@@ -1845,10 +1879,6 @@ class RebindTest(unittest.TestCase):
         ValueError, 'There are no values to rebind.'):
       Dict().rebind({})
 
-    with self.assertRaisesRegex(
-        KeyError, 'Key must be string type. Encountered 1'):
-      Dict().rebind({1: 1})
-
     with self.assertRaisesRegex(
         ValueError, 'Required value is not specified.'):
       Dict(a=1, value_spec=pg_typing.Dict([('a', pg_typing.Int())])).rebind({
@@ -1859,13 +1889,17 @@ class SerializationTest(unittest.TestCase):
   """Dedicated tests for `pg.Dict` serialization."""
 
   def test_schemaless(self):
-    sd = Dict()
+    sd = Dict({1: 2})
     sd.b = 0
     sd.c = None
     sd.a = 'foo'
 
     # Key order is preserved.
-    self.assertEqual(sd.to_json_str(), '{"b": 0, "c": null, "a": "foo"}')
+    self.assertEqual(
+        sd.to_json_str(),
+        '{"n_:1": 2, "b": 0, "c": null, "a": "foo"}'
+    )
+    self.assertEqual(base.from_json_str(sd.to_json_str()), sd)
 
   def test_schematized(self):
     sd = Dict.partial(
@@ -2076,7 +2110,7 @@ class FormatTest(unittest.TestCase):
 
   def test_compact_python_format(self):
     self.assertEqual(
-        object_utils.format(
+        utils.format(
             self._dict, compact=True, python_format=True, markdown=True
         ),
         "`{'a1': 1, 'a2': {'b1': {'c1': [{'d1': MISSING_VALUE, "
@@ -2086,9 +2120,12 @@ class FormatTest(unittest.TestCase):
 
   def test_noncompact_python_format(self):
     self.assertEqual(
-        object_utils.format(
-            self._dict, compact=False, verbose=False,
-            python_format=True, markdown=True
+        utils.format(
+            self._dict,
+            compact=False,
+            verbose=False,
+            python_format=True,
+            markdown=True,
         ),
         inspect.cleandoc("""
             ```
@@ -2188,6 +2225,7 @@ class FormatTest(unittest.TestCase):
         }"""))
 
   def test_noncompact_verbose_with_extra_blankline_for_field_docstr(self):
+    self.maxDiff = None
     self.assertEqual(
         self._dict.format(
             compact=False, verbose=True, extra_blankline_for_field_docstr=True),
@@ -2306,6 +2344,25 @@ class FormatTest(unittest.TestCase):
         '{\n  x = 1,\n  y = True,\n  z = {\n    v = 1,\n    w = True\n  }\n}'
     )
 
+  def test_compact_int_key(self):
+    d = Dict({1: 1, '1': 2})
+    self.assertEqual(d.format(compact=True), '{[1]=1, 1=2}')
+    self.assertEqual(
+        d.format(compact=True, python_format=True),
+        '{1: 1, \'1\': 2}'
+    )
+
+  def test_non_compact_int_key(self):
+    d = Dict({1: 1, '1': 2})
+    self.assertEqual(
+        d.format(compact=False, verbose=False),
+        '{\n  [1] = 1,\n  1 = 2\n}'
+    )
+    self.assertEqual(
+        d.format(compact=False, python_format=True),
+        '{\n  1: 1,\n  \'1\': 2\n}'
+    )
+
 
 def _on_change_callback(updates):
   del updates

pyglove/core/symbolic/diff.py

@@ -15,8 +15,8 @@
 
 from typing import Any, Callable, Optional, Sequence, Tuple, Union
 
-from pyglove.core import object_utils
 from pyglove.core import typing as pg_typing
+from pyglove.core import utils
 from pyglove.core.symbolic import base
 from pyglove.core.symbolic import list as pg_list
 from pyglove.core.symbolic import object as pg_object
@@ -112,8 +112,7 @@ class Diff(PureSymbolic, pg_object.Object, tree_view.HtmlTreeView.Extension):
         return 'No diff'
       # When there is no diff, but the same value needs to be displayed
       # we simply return the value.
-      return object_utils.format(
-          self.value, compact, verbose, root_indent, **kwargs)
+      return utils.format(self.value, compact, verbose, root_indent, **kwargs)
     if self.is_leaf:
       exclude_keys = kwargs.pop('exclude_keys', None)
       exclude_keys = exclude_keys or set()
@@ -129,7 +128,7 @@ class Diff(PureSymbolic, pg_object.Object, tree_view.HtmlTreeView.Extension):
             verbose=verbose,
             root_indent=root_indent,
             cls_name='',
-            bracket_type=object_utils.BracketType.SQUARE,
+            bracket_type=utils.BracketType.SQUARE,
             **kwargs,
         )
       if self.left is self.right:
@@ -141,7 +140,7 @@ class Diff(PureSymbolic, pg_object.Object, tree_view.HtmlTreeView.Extension):
           verbose=verbose,
           root_indent=root_indent,
           cls_name=cls_name,
-          bracket_type=object_utils.BracketType.ROUND,
+          bracket_type=utils.BracketType.ROUND,
           **kwargs,
       )
 
@@ -155,7 +154,7 @@ class Diff(PureSymbolic, pg_object.Object, tree_view.HtmlTreeView.Extension):
       max_summary_len_for_str: int = 80,
       **kwargs,
     ) -> Optional[tree_view.Html]:
-  # pytype: enable=annotation-type-mismatch
+    # pytype: enable=annotation-type-mismatch
     if not bool(self):
       v = self.value
       if (isinstance(v, (int, float, bool, type(None)))
@@ -199,11 +198,11 @@ class Diff(PureSymbolic, pg_object.Object, tree_view.HtmlTreeView.Extension):
       *,
       view: tree_view.HtmlTreeView,
       parent: Any = None,
-      root_path: Optional[object_utils.KeyPath] = None,
+      root_path: Optional[utils.KeyPath] = None,
       css_classes: Optional[Sequence[str]] = None,
-      **kwargs
+      **kwargs,
   ) -> tree_view.Html:
-    root_path = root_path or object_utils.KeyPath()
+    root_path = root_path or utils.KeyPath()
     if not bool(self):
       if self.value == Diff.MISSING:
         root = tree_view.Html.element(
@@ -353,7 +352,8 @@ def diff(
     right: Any,
     flatten: bool = False,
     collapse: Union[bool, str, Callable[[Any, Any], bool]] = 'same_type',
-    mode: str = 'diff') -> object_utils.Nestable[Diff]:
+    mode: str = 'diff',
+) -> utils.Nestable[Diff]:
   """Inspect the symbolic diff between two objects.
 
   For example::
@@ -479,7 +479,7 @@ def diff(
       assert isinstance(container, base.Symbolic)
       return container.sym_hasattr, container.sym_getattr, container.sym_items
 
-  def _diff(x, y) -> Tuple[object_utils.Nestable[Diff], bool]:
+  def _diff(x, y) -> Tuple[utils.Nestable[Diff], bool]:
     if x is y or x == y:
       return (Diff(x, y), False)
     if not _should_collapse(x, y):
@@ -533,5 +533,5 @@ def diff(
   if not has_diff and mode == 'diff':
     diff_value = Diff()
   if flatten:
-    diff_value = object_utils.flatten(diff_value)
+    diff_value = utils.flatten(diff_value)
   return diff_value

pyglove/core/symbolic/flags.py

@@ -14,7 +14,7 @@
 """Global, thread-local and scoped flags for handling symbolic objects."""
 
 from typing import Any, Callable, ContextManager, Optional
-from pyglove.core.object_utils import thread_local
+from pyglove.core.utils import thread_local
 
 
 #