pyglove 0.5.0.dev202510020810__py3-none-any.whl → 0.5.0.dev202512280810__py3-none-any.whl

pyglove/core/symbolic/base.py

@@ -503,10 +503,12 @@ class Symbolic(
         return default
 
   def _sym_inferred(self, key: Union[str, int], **kwargs) -> Any:
-    v = self.sym_getattr(key)
-    if isinstance(v, Inferential):
-      v = v.infer(**kwargs)
-    return v
+    return self._infer_if_applicable(self.sym_getattr(key), **kwargs)
+
+  def _infer_if_applicable(self, value: Any, **kwargs) -> Any:
+    if isinstance(value, Inferential):
+      return value.infer(**kwargs)
+    return value
 
   @abc.abstractmethod
   def sym_keys(self) -> Iterator[Union[str, int]]:
@@ -944,7 +946,7 @@ class Symbolic(
 
   def to_json(self, **kwargs) -> utils.JSONValueType:
     """Alias for `sym_jsonify`."""
-    return to_json(self, **kwargs)
+    return utils.to_json(self, **kwargs)
 
   def to_json_str(self, json_indent: Optional[int] = None, **kwargs) -> str:
     """Serializes current object into a JSON string."""
@@ -1983,10 +1985,12 @@ def is_abstract(x: Any) -> bool:
 def contains(
     x: Any,
     value: Any = None,
-    type: Optional[Union[    # pylint: disable=redefined-builtin
+    type: Union[    # pylint: disable=redefined-builtin
         Type[Any],
-        Tuple[Type[Any]]]]=None
-    ) -> bool:
+        Tuple[Type[Any], ...],
+        None,
+    ]=None,
+) -> bool:
   """Returns if a value contains values of specific type.
 
   Example::
@@ -2035,10 +2039,12 @@ def contains(
 def from_json(
     json_value: Any,
     *,
+    context: Optional[utils.JSONConversionContext] = None,
+    auto_symbolic: bool = True,
+    auto_import: bool = True,
+    convert_unknown: bool = False,
     allow_partial: bool = False,
     root_path: Optional[utils.KeyPath] = None,
-    auto_import: bool = True,
-    auto_dict: bool = False,
     value_spec: Optional[pg_typing.ValueSpec] = None,
     **kwargs,
 ) -> Any:
@@ -2059,14 +2065,23 @@ def from_json(
 
   Args:
     json_value: Input JSON value.
-    allow_partial: Whether to allow elements of the list to be partial.
-    root_path: KeyPath of loaded object in its object tree.
+    context: JSON conversion context.
+    auto_symbolic: If True, list and dict will be automatically converted to
+      `pg.List` and `pg.Dict`. Otherwise, they will be plain lists
+      and dicts.
     auto_import: If True, when a '_type' is not registered, PyGlove will
       identify its parent module and automatically import it. For example,
       if the type is 'foo.bar.A', PyGlove will try to import 'foo.bar' and
       find the class 'A' within the imported module.
-    auto_dict: If True, dict with '_type' that cannot be loaded will remain
-      as dict, with '_type' renamed to 'type_name'.
+    convert_unknown: If True, when a '_type' is not registered and cannot
+      be imported, PyGlove will create objects of:
+        - `pg.symbolic.UnknownType` for unknown types;
+        - `pg.symbolic.UnknownTypedObject` for objects of unknown types;
+        - `pg.symbolic.UnknownFunction` for unknown functions;
+        - `pg.symbolic.UnknownMethod` for unknown methods.
+      If False, TypeError will be raised.
+    allow_partial: Whether to allow elements of the list to be partial.
+    root_path: KeyPath of loaded object in its object tree.
     value_spec: The value spec for the symbolic list or dict.
     **kwargs: Allow passing through keyword arguments to from_json of specific
       types.
@@ -2082,10 +2097,23 @@ def from_json(
   if isinstance(json_value, Symbolic):
     return json_value
 
+  if context is None:
+    if (isinstance(json_value, dict) and (
+        context_node := json_value.get(utils.JSONConvertible.CONTEXT_KEY))):
+      context = utils.JSONConversionContext.from_json(
+          context_node,
+          auto_import=auto_import,
+          convert_unknown=convert_unknown,
+          **kwargs
+      )
+      json_value = json_value[utils.JSONConvertible.ROOT_VALUE_KEY]
+    else:
+      context = utils.JSONConversionContext()
+
   typename_resolved = kwargs.pop('_typename_resolved', False)
   if not typename_resolved:
     json_value = utils.json_conversion.resolve_typenames(
-        json_value, auto_import=auto_import, auto_dict=auto_dict
+        json_value, auto_import, convert_unknown
     )
 
   def _load_child(k, v):
@@ -2094,6 +2122,7 @@ def from_json(
         root_path=utils.KeyPath(k, root_path),
         _typename_resolved=True,
         allow_partial=allow_partial,
+        context=context,
         **kwargs,
     )
 
@@ -2109,24 +2138,42 @@ def from_json(
             )
         )
       return tuple(_load_child(i, v) for i, v in enumerate(json_value[1:]))
-    return Symbolic.ListType.from_json(    # pytype: disable=attribute-error
+    if json_value and json_value[0] == utils.JSONConvertible.SYMBOLIC_MARKER:
+      auto_symbolic = True
+    if auto_symbolic:
+      from_json_fn = Symbolic.ListType.from_json  # pytype: disable=attribute-error
+    else:
+      from_json_fn = utils.from_json
+    return from_json_fn(
         json_value,
+        context=context,
         value_spec=value_spec,
         root_path=root_path,
         allow_partial=allow_partial,
         **kwargs,
     )
   elif isinstance(json_value, dict):
+    if utils.JSONConvertible.REF_KEY in json_value:
+      x = context.get_shared(
+          json_value[utils.JSONConvertible.REF_KEY]
+      ).value
+      return x
     if utils.JSONConvertible.TYPE_NAME_KEY not in json_value:
-      return Symbolic.DictType.from_json(   # pytype: disable=attribute-error
-          json_value,
-          value_spec=value_spec,
-          root_path=root_path,
-          allow_partial=allow_partial,
-          **kwargs,
+      auto_symbolic = json_value.get(
+          utils.JSONConvertible.SYMBOLIC_MARKER, auto_symbolic
       )
+      if auto_symbolic:
+        return Symbolic.DictType.from_json(   # pytype: disable=attribute-error
+            json_value,
+            context=context,
+            value_spec=value_spec,
+            root_path=root_path,
+            allow_partial=allow_partial,
+            **kwargs,
+        )
     return utils.from_json(
         json_value,
+        context=context,
         _typename_resolved=True,
         root_path=root_path,
         allow_partial=allow_partial,
@@ -2138,10 +2185,12 @@ def from_json(
 def from_json_str(
     json_str: str,
     *,
+    context: Optional[utils.JSONConversionContext] = None,
+    auto_import: bool = True,
+    convert_unknown: bool = False,
     allow_partial: bool = False,
     root_path: Optional[utils.KeyPath] = None,
-    auto_import: bool = True,
-    auto_dict: bool = False,
+    value_spec: Optional[pg_typing.ValueSpec] = None,
     **kwargs,
 ) -> Any:
   """Deserialize (maybe) symbolic object from JSON string.
@@ -2161,15 +2210,22 @@ def from_json_str(
 
   Args:
     json_str: JSON string.
-    allow_partial: If True, allow a partial symbolic object to be created.
-      Otherwise error will be raised on partial value.
-    root_path: The symbolic path used for the deserialized root object.
+    context: JSON conversion context.
     auto_import: If True, when a '_type' is not registered, PyGlove will
       identify its parent module and automatically import it. For example,
       if the type is 'foo.bar.A', PyGlove will try to import 'foo.bar' and
       find the class 'A' within the imported module.
-    auto_dict: If True, dict with '_type' that cannot be loaded will remain
-      as dict, with '_type' renamed to 'type_name'.
+    convert_unknown: If True, when a '_type' is not registered and cannot
+      be imported, PyGlove will create objects of:
+        - `pg.symbolic.UnknownType` for unknown types;
+        - `pg.symbolic.UnknownTypedObject` for objects of unknown types;
+        - `pg.symbolic.UnknownFunction` for unknown functions;
+        - `pg.symbolic.UnknownMethod` for unknown methods.
+      If False, TypeError will be raised.
+    allow_partial: If True, allow a partial symbolic object to be created.
+      Otherwise error will be raised on partial value.
+    root_path: The symbolic path used for the deserialized root object.
+    value_spec: The value spec for the symbolic list or dict.
     **kwargs: Additional keyword arguments that will be passed to
       ``pg.from_json``.
 
@@ -2193,10 +2249,12 @@ def from_json_str(
 
   return from_json(
       _decode_int_keys(json.loads(json_str)),
+      context=context,
+      auto_import=auto_import,
+      convert_unknown=convert_unknown,
       allow_partial=allow_partial,
       root_path=root_path,
-      auto_import=auto_import,
-      auto_dict=auto_dict,
+      value_spec=value_spec,
       **kwargs
   )
 
@@ -2232,10 +2290,6 @@ def to_json(value: Any, **kwargs) -> Any:
   Returns:
     JSON value.
   """
-  # NOTE(daiyip): special handling `sym_jsonify` since symbolized
-  # classes may have conflicting `to_json` method in their existing classes.
-  if isinstance(value, Symbolic):
-    return value.sym_jsonify(**kwargs)
   return utils.to_json(value, **kwargs)
 
 

pyglove/core/symbolic/base_test.py

@@ -20,9 +20,9 @@ from pyglove.core import typing as pg_typing
 from pyglove.core import utils
 from pyglove.core import views
 from pyglove.core.symbolic import base
-from pyglove.core.symbolic.dict import Dict
-from pyglove.core.symbolic.inferred import ValueFromParentChain
-from pyglove.core.symbolic.object import Object
+from pyglove.core.symbolic.dict import Dict                      # pylint: disable=g-importing-member
+from pyglove.core.symbolic.inferred import ValueFromParentChain  # pylint: disable=g-importing-member
+from pyglove.core.symbolic.object import Object                  # pylint: disable=g-importing-member
 
 
 class FieldUpdateTest(unittest.TestCase):

pyglove/core/symbolic/dict.py

@@ -156,6 +156,8 @@ class Dict(dict, base.Symbolic, pg_typing.CustomTyping):
         # Not okay:
         d.a.f2.abc = 1
     """
+    # Remove symbolic marker if present.
+    json_value.pop(utils.JSONConvertible.SYMBOLIC_MARKER, None)
     return cls(
         {
             k: base.from_json(
@@ -236,7 +238,8 @@ class Dict(dict, base.Symbolic, pg_typing.CustomTyping):
         accessor_writable=True,
         # We delay seal operation until members are filled.
         sealed=False,
-        root_path=root_path)
+        root_path=root_path
+    )
 
     dict.__init__(self)
     self._value_spec = None
@@ -247,9 +250,10 @@ class Dict(dict, base.Symbolic, pg_typing.CustomTyping):
     for k, v in kwargs.items():
       dict_obj[k] = v
 
+    iter_items = getattr(dict_obj, 'sym_items', dict_obj.items)
     if value_spec:
       if pass_through:
-        for k, v in dict_obj.items():
+        for k, v in iter_items():
           super().__setitem__(k, self._relocate_if_symbolic(k, v))
 
         # NOTE(daiyip): when pass_through is on, we simply trust input
@@ -258,11 +262,11 @@ class Dict(dict, base.Symbolic, pg_typing.CustomTyping):
         # repeated validation and transformation.
         self._value_spec = value_spec
       else:
-        for k, v in dict_obj.items():
+        for k, v in iter_items():
           super().__setitem__(k, self._formalized_value(k, None, v))
         self.use_value_spec(value_spec, allow_partial)
     else:
-      for k, v in dict_obj.items():
+      for k, v in iter_items():
         self._set_item_without_permission_check(k, v)
 
     # NOTE(daiyip): We set onchange callback at the end of init to avoid
@@ -537,7 +541,7 @@ class Dict(dict, base.Symbolic, pg_typing.CustomTyping):
       raise KeyError(self._error_message(
           f'Key must be string or int type. Encountered {key!r}.'))
 
-    old_value = self.get(key, pg_typing.MISSING_VALUE)
+    old_value = self.sym_getattr(key, pg_typing.MISSING_VALUE)
     if old_value is value:
       return None
 
@@ -644,6 +648,13 @@ class Dict(dict, base.Symbolic, pg_typing.CustomTyping):
     except AttributeError as e:
       raise KeyError(key) from e
 
+  def get(self, key: Union[str, int], default: Any = None) -> Any:
+    """Get item in this Dict."""
+    try:
+      return self.sym_inferred(key)
+    except AttributeError:
+      return default
+
   def __setitem__(self, key: Union[str, int], value: Any) -> None:
     """Set item in this Dict.
 
@@ -751,11 +762,13 @@ class Dict(dict, base.Symbolic, pg_typing.CustomTyping):
 
   def items(self) -> Iterator[Tuple[Union[str, int], Any]]:  # pytype: disable=signature-mismatch
     """Returns an iterator of (key, value) items in current dict."""
-    return self.sym_items()
+    for k, v in self.sym_items():
+      yield k, self._infer_if_applicable(v)
 
   def values(self) -> Iterator[Any]:  # pytype: disable=signature-mismatch
     """Returns an iterator of values in current dict.."""
-    return self.sym_values()
+    for v in self.sym_values():
+      yield self._infer_if_applicable(v)
 
   def copy(self) -> 'Dict':
     """Overridden copy using symbolic copy."""
@@ -824,12 +837,15 @@ class Dict(dict, base.Symbolic, pg_typing.CustomTyping):
       hide_default_values: bool = False,
       exclude_keys: Optional[Sequence[Union[str, int]]] = None,
       use_inferred: bool = False,
+      omit_symbolic_marker: bool = True,
       **kwargs,
   ) -> utils.JSONValueType:
     """Converts current object to a dict with plain Python objects."""
     exclude_keys = set(exclude_keys or [])
+    json_repr = {}
+    if not omit_symbolic_marker:
+      json_repr[utils.JSONConvertible.SYMBOLIC_MARKER] = True
     if self._value_spec and self._value_spec.schema:
-      json_repr = dict()
       matched_keys, _ = self._value_spec.schema.resolve(self.keys())  # pytype: disable=attribute-error
       for key_spec, keys in matched_keys.items():
         # NOTE(daiyip): The key values of frozen field can safely be excluded
@@ -851,20 +867,23 @@ class Dict(dict, base.Symbolic, pg_typing.CustomTyping):
                 hide_frozen=hide_frozen,
                 hide_default_values=hide_default_values,
                 use_inferred=use_inferred,
-                **kwargs)
-      return json_repr
+                omit_symbolic_marker=omit_symbolic_marker,
+                **kwargs
+            )
     else:
-      return {
+      json_repr.update({
           k: base.to_json(
               self.sym_inferred(k, default=v) if (
                   use_inferred and isinstance(v, base.Inferential)) else v,
               hide_frozen=hide_frozen,
               hide_default_values=hide_default_values,
               use_inferred=use_inferred,
+              omit_symbolic_marker=omit_symbolic_marker,
               **kwargs)
           for k, v in self.sym_items()
           if k not in exclude_keys
-      }
+      })
+    return json_repr
 
   def custom_apply(
       self,

pyglove/core/symbolic/dict_test.py

@@ -415,6 +415,19 @@ class DictTest(unittest.TestCase):
     with self.assertRaisesRegex(KeyError, 'Key \'y1\' is not allowed'):
       sd.y1 = 4
 
+  def test_get(self):
+    sd = Dict(a=1)
+    self.assertEqual(sd.get('a'), 1)
+    self.assertIsNone(sd.get('x'))
+    self.assertEqual(sd.get('x', 2), 2)
+
+    # Test inferred values.
+    sd = Dict(x=inferred.ValueFromParentChain())
+    self.assertIsNone(sd.get('x'))
+
+    _ = Dict(sd=sd, x=1)
+    self.assertEqual(sd.get('x'), 1)
+
   def test_getattr(self):
     sd = Dict(a=1)
     self.assertEqual(sd.a, 1)
@@ -713,6 +726,11 @@ class DictTest(unittest.TestCase):
     ]))
     self.assertEqual(list(sd.values()), [2, 1, 3])
 
+    # Test values with inferred values.
+    sd = Dict(x=1, y=inferred.ValueFromParentChain())
+    _ = Dict(sd=sd, y=2)
+    self.assertEqual(list(sd.values()), [1, 2])
+
   def test_items(self):
     sd = Dict(b={'c': True, 'd': []}, a=0)
     self.assertEqual(list(sd.items()), [('b', {'c': True, 'd': []}), ('a', 0)])
@@ -725,6 +743,11 @@ class DictTest(unittest.TestCase):
     ]))
     self.assertEqual(list(sd.items()), [('b', 2), ('a', 1), ('c', 3)])
 
+    # Test items with inferred values.
+    sd = Dict(x=1, y=inferred.ValueFromParentChain())
+    _ = Dict(sd=sd, y=2)
+    self.assertEqual(list(sd.items()), [('x', 1), ('y', 2)])
+
   def test_non_default(self):
     sd = Dict(a=1)
     self.assertEqual(len(sd.non_default_values()), 1)
@@ -923,6 +946,14 @@ class DictTest(unittest.TestCase):
         sd.sym_jsonify(),
         {'x': 1, 'y': inferred.ValueFromParentChain().to_json()},
     )
+    self.assertEqual(
+        sd.sym_jsonify(omit_symbolic_marker=False),
+        {
+            '__symbolic__': True,
+            'x': 1,
+            'y': inferred.ValueFromParentChain().to_json()
+        },
+    )
 
   def test_sym_rebind(self):
     # Refer to RebindTest for more detailed tests.
@@ -1941,6 +1972,24 @@ class SerializationTest(unittest.TestCase):
     self.assertEqual(sd.to_json_str(), '{"x": 1, "y": 2.0}')
     self.assertEqual(base.from_json_str(sd.to_json_str(), value_spec=spec), sd)
 
+  def test_auto_symbolic(self):
+    value = base.from_json({'x': 1}, auto_symbolic=True)
+    self.assertIsInstance(value, Dict)
+
+    value = base.from_json({'x': 1}, auto_symbolic=False)
+    self.assertNotIsInstance(value, Dict)
+
+  def test_omit_symbolic_marker(self):
+    sd = Dict(x=1)
+    self.assertEqual(sd.to_json(omit_symbolic_marker=True), {'x': 1})
+    self.assertEqual(
+        sd.to_json(omit_symbolic_marker=False),
+        {'__symbolic__': True, 'x': 1}
+    )
+    sd = base.from_json({'__symbolic__': True, 'x': 1}, auto_symbolic=False)
+    self.assertIsInstance(sd, Dict)
+    self.assertEqual(sd, {'x': 1})
+
   def test_hide_frozen(self):
 
     class A(pg_object.Object):

pyglove/core/symbolic/list.py

@@ -137,6 +137,9 @@ class List(list, base.Symbolic, pg_typing.CustomTyping):
     Returns:
       A schema-less symbolic list, but its items maybe symbolic.
     """
+    # Remove symbolic marker if present.
+    if json_value and json_value[0] == utils.JSONConvertible.SYMBOLIC_MARKER:
+      json_value.pop(0)
     return cls(
         [
             base.from_json(
@@ -770,15 +773,26 @@ class List(list, base.Symbolic, pg_typing.CustomTyping):
     return (proceed_with_standard_apply, self)
 
   def sym_jsonify(
-      self, use_inferred: bool = False, **kwargs
+      self,
+      use_inferred: bool = False,
+      omit_symbolic_marker: bool = True,
+      **kwargs
   ) -> utils.JSONValueType:
     """Converts current list to a list of plain Python objects."""
     def json_item(idx):
       v = self.sym_getattr(idx)
       if use_inferred and isinstance(v, base.Inferential):
         v = self.sym_inferred(idx, default=v)
-      return base.to_json(v, use_inferred=use_inferred, **kwargs)
-    return [json_item(i) for i in range(len(self))]
+      return base.to_json(
+          v,
+          use_inferred=use_inferred,
+          omit_symbolic_marker=omit_symbolic_marker,
+          **kwargs
+      )
+    json_value = [json_item(i) for i in range(len(self))]
+    if not omit_symbolic_marker:
+      json_value.insert(0, utils.JSONConvertible.SYMBOLIC_MARKER)
+    return json_value
 
   def format(
       self,

pyglove/core/symbolic/list_test.py

@@ -506,8 +506,7 @@ class ListTest(unittest.TestCase):
   def test_index(self):
     sl = List([0, 1, 2, 1])
     self.assertEqual(sl.index(1), 1)
-    with self.assertRaisesRegex(
-        ValueError, '3 is not in list'):
+    with self.assertRaisesRegex(ValueError, '.* not in list'):
       _ = sl.index(3)
 
     # Index of inferred value is based on its symbolic form.
@@ -799,6 +798,10 @@ class ListTest(unittest.TestCase):
     self.assertEqual(
         sl.sym_jsonify(), [0, inferred.ValueFromParentChain().to_json()]
     )
+    self.assertEqual(
+        sl.sym_jsonify(omit_symbolic_marker=False),
+        ['__symbolic__', 0, inferred.ValueFromParentChain().to_json()]
+    )
 
   def test_sym_rebind(self):
     # Refer to RebindTest for more detailed tests.
@@ -1633,6 +1636,25 @@ class SerializationTest(unittest.TestCase):
         'Tuple should have at least one element besides \'__tuple__\'.'):
       base.from_json_str('["__tuple__"]')
 
+  def test_auto_symbolic(self):
+    value = base.from_json([1, 2, 3], auto_symbolic=True)
+    self.assertIsInstance(value, List)
+
+    value = base.from_json([1, 2, 3], auto_symbolic=False)
+    self.assertNotIsInstance(value, List)
+
+  def test_omit_symbolic_marker(self):
+    sl = List([0])
+    self.assertEqual(sl.to_json(omit_symbolic_marker=True), [0])
+    self.assertEqual(
+        sl.to_json(omit_symbolic_marker=False),
+        ['__symbolic__', 0]
+    )
+    sl = base.from_json(['__symbolic__', 0], auto_symbolic=False)
+    self.assertEqual(sl.to_json(omit_symbolic_marker=True), [0])
+    self.assertIsInstance(sl, List)
+    self.assertEqual(sl, [0])
+
   def test_hide_default_values(self):
     sl = List.partial(
         [dict(x=1)],

pyglove/core/symbolic/object.py

@@ -339,7 +339,8 @@ class Object(base.Symbolic, metaclass=ObjectMeta):
 
     # Set `__serialization_key__` before JSONConvertible.__init_subclass__
     # is called.
-    setattr(cls, '__serialization_key__', cls.__type_name__)
+    if '__serialization_key__' not in cls.__dict__:
+      setattr(cls, '__serialization_key__', cls.__type_name__)
 
     super().__init_subclass__()
 
@@ -994,6 +995,7 @@ class Object(base.Symbolic, metaclass=ObjectMeta):
             self.__class__.__serialization_key__
         )
     }
+    kwargs['omit_symbolic_marker'] = True
     json_dict.update(self._sym_attributes.to_json(**kwargs))
     return json_dict
 

pyglove/core/symbolic/object_test.py

@@ -38,6 +38,7 @@ from pyglove.core.symbolic.object import use_init_args as pg_use_init_args
 from pyglove.core.symbolic.origin import Origin
 from pyglove.core.symbolic.pure_symbolic import NonDeterministic
 from pyglove.core.symbolic.pure_symbolic import PureSymbolic
+from pyglove.core.symbolic.unknown_symbols import UnknownTypedObject
 from pyglove.core.views.html import tree_view  # pylint: disable=unused-import
 
 
@@ -3158,7 +3159,7 @@ class SerializationTest(unittest.TestCase):
             Q.partial(P.partial()).to_json_str(), allow_partial=True),
         Q.partial(P.partial()))
 
-  def test_serialization_with_auto_dict(self):
+  def test_serialization_with_convert_unknown(self):
 
     class P(Object):
       auto_register = False
@@ -3181,15 +3182,17 @@ class SerializationTest(unittest.TestCase):
         }
     )
     self.assertEqual(
-        base.from_json_str(Q(P(1), y='foo').to_json_str(), auto_dict=True),
-        {
-            'p': {
-                'type_name': P.__type_name__,
-                'x': 1
-            },
-            'y': 'foo',
-            'type_name': Q.__type_name__,
-        }
+        base.from_json_str(
+            Q(P(1), y='foo').to_json_str(), convert_unknown=True
+        ),
+        UnknownTypedObject(
+            type_name=Q.__type_name__,
+            p=UnknownTypedObject(
+                type_name=P.__type_name__,
+                x=1
+            ),
+            y='foo'
+        )
     )
 
   def test_serialization_with_converter(self):

pyglove/core/symbolic/ref.py

@@ -138,9 +138,7 @@ class Ref(
     del child_transform
     # Check if the field being assigned could accept the referenced value.
     # We do not do any transformation, thus not passing the child transform.
-    value_spec.apply(
-        self._value,
-        allow_partial=allow_partial)
+    value_spec.apply(self._value, allow_partial=allow_partial)
     return (False, self)
 
   def _sym_clone(self, deep: bool, memo: Any = None) -> 'Ref':
@@ -152,10 +150,24 @@ class Ref(
   def sym_eq(self, other: Any) -> bool:
     return isinstance(other, Ref) and self.value is other.value
 
-  def sym_jsonify(self, *, save_ref_value: bool = False, **kwargs: Any) -> Any:
-    if save_ref_value:
-      return base.to_json(self._value, save_ref_value=save_ref_value, **kwargs)
-    raise TypeError(f'{self!r} cannot be serialized at the moment.')
+  def sym_jsonify(
+      self,
+      *,
+      context: utils.JSONConversionContext,
+      **kwargs: Any
+  ) -> Any:
+    # Disable auto_symbolic for Ref value. This allows Ref to create a sub-tree
+    # for reference sharing.
+    kwargs['omit_symbolic_marker'] = False
+    return {
+        utils.JSONConvertible.TYPE_NAME_KEY: self.__class__.__type_name__,
+        'value': context.serialize_maybe_shared(self._value, **kwargs)
+    }
+
+  @classmethod
+  def from_json(cls, json: Any, **kwargs):
+    kwargs['auto_symbolic'] = False
+    return super().from_json(json, **kwargs)
 
   def __getstate__(self):
     raise TypeError(f'{self!r} cannot be pickled at the moment.')