Add Symbolic.sym_root property.

This CL allows convenient access to the root of current symbolic tree.

PiperOrigin-RevId: 516366708

docs/learn/soop/som/operations.rst

@@ -41,7 +41,7 @@ Property :attr:`sym_parent <pyglove.symbolic.Symbolic.sym_parent>` is the API fo
 symbolic types to access their containing node (parent) in the tree. For example, the ``Cage``
 object in the `zoo` has ``exhibits`` (a symbolic list) as its parent::
 
-    assert zoo.exhibits[0].sym_parent == zoo.exhibits[0]
+    assert zoo.exhibits[0].sym_parent is zoo.exhibits[0]
 
 
 .. tip::
@@ -57,6 +57,16 @@ object in the `zoo` has ``exhibits`` (a symbolic list) as its parent::
          assert shark.sym_parent is pool
          assert shark.sym_path == 'x'
 
+Root
+====
+
+Similarly, users can access the root of current symbolic tree via
+property :attr:`sym_root <pyglove.symbolic.Symbolic.sym_root>`. For example::
+
+  assert zoo.sym_root is zoo
+  assert zoo.exhibits[0].sym_root is zoo
+  assert zoo.exhibits[0].animal.sym_root is zoo
+
 
 Child Nodes
 ===========

pyglove/core/symbolic/base.py

@@ -238,6 +238,14 @@ def sym_field(self) -> Optional[pg_typing.Field]:
       return None
     return self.sym_parent.sym_attr_field(self.sym_path.key)
 
+  @property
+  def sym_root(self) -> 'Symbolic':
+    """Returns the root of the symbolic tree."""
+    root = self
+    while root.sym_parent is not None:
+      root = root.sym_parent
+    return root
+
   @abc.abstractmethod
   def sym_attr_field(self, key: Union[str, int]) -> Optional[pg_typing.Field]:
     """Returns the field definition for a symbolic attribute."""

pyglove/core/symbolic/dict_test.py

@@ -1104,6 +1104,15 @@ def test_sym_parent(self):
     pd = Dict(a=sd)
     self.assertIs(sd.sym_parent, pd)
 
+  def test_sym_root(self):
+    sd = Dict(x=dict(a=1), y=[])
+    self.assertIs(sd.sym_root, sd)
+    self.assertIs(sd.x.sym_parent, sd)
+    self.assertIs(sd.y.sym_parent, sd)
+
+    pd = Dict(a=sd)
+    self.assertIs(sd.sym_root, pd)
+
   def test_sym_path(self):
     sd = Dict(x=dict(a=dict()), y=[dict(b=dict())])
     self.assertEqual(sd.sym_path, '')

pyglove/core/symbolic/list_test.py

@@ -958,6 +958,18 @@ def test_sym_parent(self):
     pl = List([sl])
     self.assertIs(sl.sym_parent, pl)
 
+  def test_sym_root(self):
+    sl = List([[0], dict(x=1)])
+    self.assertIs(sl.sym_root, sl)
+
+    self.assertIs(sl[0].sym_root, sl)
+    self.assertIs(sl[1].sym_root, sl)
+
+    pl = List([sl])
+    self.assertIs(sl.sym_root, pl)
+    self.assertIs(sl[0].sym_root, pl)
+    self.assertIs(sl[1].sym_root, pl)
+
   def test_sym_path(self):
     sl = List([dict(a=dict()), [dict(b=[0])]])
     self.assertEqual(sl.sym_path, '')

pyglove/core/symbolic/object_test.py

@@ -1156,6 +1156,29 @@ class A(Object):
     pa = A(a)
     self.assertIs(a.sym_parent, pa)
 
+  def test_sym_root(self):
+
+    @pg_members([
+        ('x', pg_typing.Any()),
+    ])
+    class A(Object):
+      pass
+
+    a = A(dict(x=A([A(1)])))
+    self.assertIs(a.sym_root, a)
+
+    self.assertIs(a.x.sym_root, a)
+    self.assertIs(a.x.x.sym_root, a)
+    self.assertIs(a.x.x.x.sym_root, a)
+    self.assertIs(a.x.x.x[0].sym_root, a)
+
+    pa = A(a)
+    self.assertIs(a.sym_root, pa)
+    self.assertIs(a.x.sym_root, pa)
+    self.assertIs(a.x.x.sym_root, pa)
+    self.assertIs(a.x.x.x.sym_root, pa)
+    self.assertIs(a.x.x.x[0].sym_root, pa)
+
   def test_sym_path(self):
 
     @pg_members([