pyglove 0.4.5.dev20240319__py3-none-any.whl → 0.4.5.dev202501132210__py3-none-any.whl

pyglove/core/geno/numerical.py

@@ -17,10 +17,9 @@ import random
 import types
 from typing import Any, List, Optional, Union
 
-from pyglove.core import object_utils
 from pyglove.core import symbolic
 from pyglove.core import typing as pg_typing
-
+from pyglove.core import utils
 from pyglove.core.geno.base import DecisionPoint
 from pyglove.core.geno.base import DNA
 
@@ -164,25 +163,30 @@ class Float(DecisionPoint):
     if not compact:
       return super().format(compact, verbose, root_indent, **kwargs)
     if show_id:
-      kvlist = [('id', object_utils.quote_if_str(str(self.id)), '\'\'')]
+      kvlist = [('id', str(self.id), '\'\'')]
     else:
       kvlist = []
-    details = object_utils.kvlist_str(kvlist + [
-        ('name', object_utils.quote_if_str(self.name), None),
-        ('min_value', self.min_value, None),
-        ('max_value', self.max_value, None),
-        ('scale', self.scale, None),
-        ('hints', object_utils.quote_if_str(self.hints), None),
-    ])
+    details = utils.kvlist_str(
+        kvlist
+        + [
+            ('name', self.name, None),
+            ('min_value', self.min_value, None),
+            ('max_value', self.max_value, None),
+            ('scale', self.scale, None),
+            ('hints', self.hints, None),
+        ]
+    )
     return f'{self.__class__.__name__}({details})'
 
 
-def floatv(min_value: float,
-           max_value: float,
-           scale: Optional[str] = None,
-           hints: Any = None,
-           location: object_utils.KeyPath = object_utils.KeyPath(),
-           name: Optional[str] = None) -> Float:
+def floatv(
+    min_value: float,
+    max_value: float,
+    scale: Optional[str] = None,
+    hints: Any = None,
+    location: utils.KeyPath = utils.KeyPath(),
+    name: Optional[str] = None,
+) -> Float:
   """Returns a Float specification.
 
   It creates the genotype for :func:`pyglove.floatv`.
@@ -226,4 +230,3 @@ def floatv(min_value: float,
   """
   return Float(min_value, max_value, scale,
                hints=hints, location=location, name=name)
-

pyglove/core/geno/space.py

@@ -18,10 +18,9 @@ import random
 import types
 from typing import List, Optional, Union
 
-from pyglove.core import object_utils
 from pyglove.core import symbolic
 from pyglove.core import typing as pg_typing
-
+from pyglove.core import utils
 from pyglove.core.geno.base import DecisionPoint
 from pyglove.core.geno.base import DNA
 from pyglove.core.geno.base import DNASpec
@@ -138,7 +137,7 @@ class Space(DNASpec):
     """
     return not self.elements
 
-  def validate(self, dna: DNA) -> None:
+  def validate(self, dna: DNA) -> None:  # pytype: disable=signature-mismatch
     """Validate whether a DNA value conforms to this spec."""
     if not self.elements and (dna.value is not None or dna.children):
       raise ValueError(
@@ -231,8 +230,8 @@ class Space(DNASpec):
     return sum([len(elem) for elem in self.elements])
 
   def __getitem__(
-      self, index: Union[int, slice, str, object_utils.KeyPath]
-      ) -> Union[DecisionPoint, List[DecisionPoint]]:
+      self, index: Union[int, slice, str, utils.KeyPath]
+  ) -> Union[DecisionPoint, List[DecisionPoint]]:
     """Operator [] to return element by index or sub-DNASpec by name."""
     if isinstance(index, (int, slice)):
       return self.elements[index]

pyglove/core/hyper/base.py

@@ -18,9 +18,9 @@ import abc
 from typing import Any, Callable, Optional
 
 from pyglove.core import geno
-from pyglove.core import object_utils
 from pyglove.core import symbolic
 from pyglove.core import typing as pg_typing
+from pyglove.core import utils
 
 
 class HyperValue(symbolic.NonDeterministic):  # pytype: disable=ignored-metaclass
@@ -110,8 +110,7 @@ class HyperValue(symbolic.NonDeterministic):  # pytype: disable=ignored-metaclas
     """
 
   @abc.abstractmethod
-  def dna_spec(self,
-               location: Optional[object_utils.KeyPath] = None) -> geno.DNASpec:
+  def dna_spec(self, location: Optional[utils.KeyPath] = None) -> geno.DNASpec:
     """Get DNA spec of DNA that is decodable/encodable by this hyper value."""
 
 
@@ -186,12 +185,13 @@ def set_dynamic_evaluate_fn(
   global _global_dynamic_evaluate_fn
   if per_thread:
     assert _global_dynamic_evaluate_fn is None, _global_dynamic_evaluate_fn
-    object_utils.thread_local_set(_TLS_KEY_DYNAMIC_EVALUATE_FN, fn)
+    utils.thread_local_set(_TLS_KEY_DYNAMIC_EVALUATE_FN, fn)
   else:
     _global_dynamic_evaluate_fn = fn
 
 
 def get_dynamic_evaluate_fn() -> Optional[Callable[[HyperValue], Any]]:
   """Gets current dynamic evaluate function."""
-  return object_utils.thread_local_get(
-      _TLS_KEY_DYNAMIC_EVALUATE_FN, _global_dynamic_evaluate_fn)
+  return utils.thread_local_get(
+      _TLS_KEY_DYNAMIC_EVALUATE_FN, _global_dynamic_evaluate_fn
+  )

pyglove/core/hyper/categorical.py

@@ -18,9 +18,9 @@ import typing
 from typing import Any, Callable, Iterable, List, Optional, Tuple, Union
 
 from pyglove.core import geno
-from pyglove.core import object_utils
 from pyglove.core import symbolic
 from pyglove.core import typing as pg_typing
+from pyglove.core import utils
 from pyglove.core.hyper import base
 from pyglove.core.hyper import object_template
 
@@ -85,7 +85,8 @@ class Choices(base.HyperPrimitive):
     self._value_spec = None
 
   def _update_children_paths(
-      self, old_path: object_utils.KeyPath, new_path: object_utils.KeyPath):
+      self, old_path: utils.KeyPath, new_path: utils.KeyPath
+  ):
     """Customized logic to update children paths."""
     super()._update_children_paths(old_path, new_path)
     for t in self._candidate_templates:
@@ -104,19 +105,20 @@ class Choices(base.HyperPrimitive):
         return False
     return True
 
-  def dna_spec(self,
-               location: Optional[object_utils.KeyPath] = None) -> geno.Choices:
+  def dna_spec(self, location: Optional[utils.KeyPath] = None) -> geno.Choices:
     """Returns corresponding DNASpec."""
     return geno.Choices(
         num_choices=self.num_choices,
         candidates=[ct.dna_spec() for ct in self._candidate_templates],
-        literal_values=[self._literal_value(c)
-                        for i, c in enumerate(self.candidates)],
+        literal_values=[
+            self._literal_value(c) for i, c in enumerate(self.candidates)
+        ],
         distinct=self.choices_distinct,
         sorted=self.choices_sorted,
         hints=self.hints,
         name=self.name,
-        location=location or object_utils.KeyPath())
+        location=location or utils.KeyPath(),
+    )
 
   def _literal_value(
       self, candidate: Any, max_len: int = 120) -> Union[int, float, str]:
@@ -124,10 +126,13 @@ class Choices(base.HyperPrimitive):
     if isinstance(candidate, numbers.Number):
       return candidate
 
-    literal = object_utils.format(candidate, compact=True,
-                                  hide_default_values=True,
-                                  hide_missing_values=True,
-                                  strip_object_id=True)
+    literal = utils.format(
+        candidate,
+        compact=True,
+        hide_default_values=True,
+        hide_missing_values=True,
+        strip_object_id=True,
+    )
     if len(literal) > max_len:
       literal = literal[:max_len - 3] + '...'
     return literal
@@ -139,52 +144,70 @@ class Choices(base.HyperPrimitive):
       # Single choice.
       if not isinstance(dna.value, int):
         raise ValueError(
-            object_utils.message_on_path(
-                f'Did you forget to specify values for conditional choices?\n'
+            utils.message_on_path(
+                'Did you forget to specify values for conditional choices?\n'
                 f'Expect integer for {self.__class__.__name__}. '
-                f'Encountered: {dna!r}.', self.sym_path))
+                f'Encountered: {dna!r}.',
+                self.sym_path,
+            )
+        )
       if dna.value >= len(self.candidates):
         raise ValueError(
-            object_utils.message_on_path(
+            utils.message_on_path(
                 f'Choice out of range. Value: {dna.value!r}, '
-                f'Candidates: {len(self.candidates)}.', self.sym_path))
+                f'Candidates: {len(self.candidates)}.',
+                self.sym_path,
+            )
+        )
       choices = [self._candidate_templates[dna.value].decode(
           geno.DNA(None, dna.children))]
     else:
       # Multi choices.
       if len(dna.children) != self.num_choices:
         raise ValueError(
-            object_utils.message_on_path(
-                f'Number of DNA child values does not match the number of '
+            utils.message_on_path(
+                'Number of DNA child values does not match the number of '
                 f'choices. Child values: {dna.children!r}, '
-                f'Choices: {self.num_choices}.', self.sym_path))
+                f'Choices: {self.num_choices}.',
+                self.sym_path,
+            )
+        )
       if self.choices_distinct or self.choices_sorted:
         sub_dna_values = [s.value for s in dna]
         if (self.choices_distinct
             and len(set(sub_dna_values)) != len(dna.children)):
           raise ValueError(
-              object_utils.message_on_path(
-                  f'DNA child values should be distinct. '
-                  f'Encountered: {sub_dna_values}.', self.sym_path))
+              utils.message_on_path(
+                  'DNA child values should be distinct. '
+                  f'Encountered: {sub_dna_values}.',
+                  self.sym_path,
+              )
+          )
         if self.choices_sorted and sorted(sub_dna_values) != sub_dna_values:
           raise ValueError(
-              object_utils.message_on_path(
-                  f'DNA child values should be sorted. '
-                  f'Encountered: {sub_dna_values}.', self.sym_path))
+              utils.message_on_path(
+                  'DNA child values should be sorted. '
+                  f'Encountered: {sub_dna_values}.',
+                  self.sym_path,
+              )
+          )
       choices = []
       for i, sub_dna in enumerate(dna):
         if not isinstance(sub_dna.value, int):
           raise ValueError(
-              object_utils.message_on_path(
-                  f'Choice value should be int. '
-                  f'Encountered: {sub_dna.value}.',
-                  object_utils.KeyPath(i, self.sym_path)))
+              utils.message_on_path(
+                  f'Choice value should be int. Encountered: {sub_dna.value}.',
+                  utils.KeyPath(i, self.sym_path),
+              )
+          )
         if sub_dna.value >= len(self.candidates):
           raise ValueError(
-              object_utils.message_on_path(
+              utils.message_on_path(
                   f'Choice out of range. Value: {sub_dna.value}, '
                   f'Candidates: {len(self.candidates)}.',
-                  object_utils.KeyPath(i, self.sym_path)))
+                  utils.KeyPath(i, self.sym_path),
+              )
+          )
         choices.append(self._candidate_templates[sub_dna.value].decode(
             geno.DNA(None, sub_dna.children)))
     return choices
@@ -240,15 +263,21 @@ class Choices(base.HyperPrimitive):
     """
     if not isinstance(value, list):
       raise ValueError(
-          object_utils.message_on_path(
-              f'Cannot encode value: value should be a list type. '
-              f'Encountered: {value!r}.', self.sym_path))
+          utils.message_on_path(
+              'Cannot encode value: value should be a list type. '
+              f'Encountered: {value!r}.',
+              self.sym_path,
+          )
+      )
     choices = []
     if self.num_choices is not None and len(value) != self.num_choices:
       raise ValueError(
-          object_utils.message_on_path(
-              f'Length of input list is different from the number of choices '
-              f'({self.num_choices}). Encountered: {value}.', self.sym_path))
+          utils.message_on_path(
+              'Length of input list is different from the number of choices '
+              f'({self.num_choices}). Encountered: {value}.',
+              self.sym_path,
+          )
+      )
     for v in value:
       choice_id = None
       child_dna = None
@@ -259,10 +288,12 @@ class Choices(base.HyperPrimitive):
           break
       if child_dna is None:
         raise ValueError(
-            object_utils.message_on_path(
-                f'Cannot encode value: no candidates matches with '
+            utils.message_on_path(
+                'Cannot encode value: no candidates matches with '
                 f'the value. Value: {v!r}, Candidates: {self.candidates}.',
-                self.sym_path))
+                self.sym_path,
+            )
+        )
       choices.append(geno.DNA(choice_id, [child_dna]))
     return geno.DNA(None, choices)
 
@@ -313,12 +344,13 @@ class ManyOf(Choices):
 
   def custom_apply(
       self,
-      path: object_utils.KeyPath,
+      path: utils.KeyPath,
       value_spec: pg_typing.ValueSpec,
       allow_partial: bool,
-      child_transform: Optional[Callable[
-          [object_utils.KeyPath, pg_typing.Field, Any], Any]] = None
-      ) -> Tuple[bool, 'Choices']:
+      child_transform: Optional[
+          Callable[[utils.KeyPath, pg_typing.Field, Any], Any]
+      ] = None,
+  ) -> Tuple[bool, 'Choices']:
     """Validate candidates during value_spec binding time."""
     # Check if value_spec directly accepts `self`.
     if value_spec.value_type and isinstance(self, value_spec.value_type):
@@ -329,10 +361,12 @@ class ManyOf(Choices):
       dest_spec = value_spec
       if not dest_spec.is_compatible(src_spec):
         raise TypeError(
-            object_utils.message_on_path(
-                f'Cannot bind an incompatible value spec {dest_spec} '
-                f'to {self.__class__.__name__} with bound spec {src_spec}.',
-                path))
+            utils.message_on_path(
+                f'Cannot bind an incompatible value spec {dest_spec!r} '
+                f'to {self.__class__.__name__} with bound spec {src_spec!r}.',
+                path,
+            )
+        )
       return (False, self)
 
     list_spec = typing.cast(
@@ -399,12 +433,13 @@ class OneOf(Choices):
 
   def custom_apply(
       self,
-      path: object_utils.KeyPath,
+      path: utils.KeyPath,
       value_spec: pg_typing.ValueSpec,
       allow_partial: bool,
-      child_transform: Optional[Callable[
-          [object_utils.KeyPath, pg_typing.Field, Any], Any]] = None
-      ) -> Tuple[bool, 'OneOf']:
+      child_transform: Optional[
+          Callable[[utils.KeyPath, pg_typing.Field, Any], Any]
+      ] = None,
+  ) -> Tuple[bool, 'OneOf']:
     """Validate candidates during value_spec binding time."""
     # Check if value_spec directly accepts `self`.
     if value_spec.value_type and isinstance(self, value_spec.value_type):
@@ -413,10 +448,13 @@ class OneOf(Choices):
     if self._value_spec:
       if not value_spec.is_compatible(self._value_spec):
         raise TypeError(
-            object_utils.message_on_path(
-                f'Cannot bind an incompatible value spec {value_spec} '
+            utils.message_on_path(
+                f'Cannot bind an incompatible value spec {value_spec!r} '
                 f'to {self.__class__.__name__} with bound '
-                f'spec {self._value_spec}.', path))
+                f'spec {self._value_spec!r}.',
+                path,
+            )
+        )
       return (False, self)
 
     for i, c in enumerate(self.candidates):
@@ -427,6 +465,7 @@ class OneOf(Choices):
     self._value_spec = value_spec
     return (False, self)
 
+
 #
 # Helper methods for creating hyper values.
 #

pyglove/core/hyper/custom.py

@@ -19,8 +19,8 @@ import types
 from typing import Any, Callable, Optional, Tuple, Union
 
 from pyglove.core import geno
-from pyglove.core import object_utils
 from pyglove.core import typing as pg_typing
+from pyglove.core import utils
 from pyglove.core.hyper import base
 
 
@@ -111,8 +111,7 @@ class CustomHyper(base.HyperPrimitive):
     raise NotImplementedError(
         f'\'custom_encode\' is not supported by {self.__class__.__name__!r}.')
 
-  def dna_spec(
-      self, location: Optional[object_utils.KeyPath] = None) -> geno.DNASpec:
+  def dna_spec(self, location: Optional[utils.KeyPath] = None) -> geno.DNASpec:
     """Always returns CustomDecisionPoint for CustomHyper."""
     return geno.CustomDecisionPoint(
         hyper_type=self.__class__.__name__,
@@ -147,12 +146,13 @@ class CustomHyper(base.HyperPrimitive):
 
   def custom_apply(
       self,
-      path: object_utils.KeyPath,
+      path: utils.KeyPath,
       value_spec: pg_typing.ValueSpec,
       allow_partial: bool,
-      child_transform: Optional[Callable[
-          [object_utils.KeyPath, pg_typing.Field, Any], Any]] = None
-      ) -> Tuple[bool, 'CustomHyper']:
+      child_transform: Optional[
+          Callable[[utils.KeyPath, pg_typing.Field, Any], Any]
+      ] = None,
+  ) -> Tuple[bool, 'CustomHyper']:
     """Validate candidates during value_spec binding time."""
     del path, value_spec, allow_partial, child_transform
     # Allow custom hyper to be assigned to any type.

pyglove/core/hyper/custom_test.py

@@ -11,15 +11,12 @@
 # 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.hyper.CustomHyper."""
-
 import random
 import unittest
 
 from pyglove.core import geno
-from pyglove.core import object_utils
 from pyglove.core import symbolic
-
+from pyglove.core import utils
 from pyglove.core.hyper.categorical import oneof
 from pyglove.core.hyper.custom import CustomHyper
 from pyglove.core.hyper.iter import iterate
@@ -58,12 +55,14 @@ class CustomHyperTest(unittest.TestCase):
   """Test for CustomHyper."""
 
   def test_dna_spec(self):
-    self.assertTrue(symbolic.eq(
-        IntSequence(hints='x').dna_spec('a'),
-        geno.CustomDecisionPoint(
-            hyper_type='IntSequence',
-            location=object_utils.KeyPath('a'),
-            hints='x')))
+    self.assertTrue(
+        symbolic.eq(
+            IntSequence(hints='x').dna_spec('a'),
+            geno.CustomDecisionPoint(
+                hyper_type='IntSequence', location=utils.KeyPath('a'), hints='x'
+            ),
+        )
+    )
 
   def test_decode(self):
     self.assertEqual(IntSequence().decode(geno.DNA('0,1,2')), [0, 1, 2])

pyglove/core/hyper/derived.py

@@ -17,16 +17,19 @@ import abc
 import copy
 from typing import Any, Callable, List, Optional, Tuple, Union
 
-from pyglove.core import object_utils
 from pyglove.core import symbolic
 from pyglove.core import typing as pg_typing
+from pyglove.core import utils
 
 
-@symbolic.members([
-    ('reference_paths', pg_typing.List(pg_typing.Object(object_utils.KeyPath)),
-     ('Paths of referenced values, which are relative paths searched from '
-      'current node to root.'))
-])
+@symbolic.members([(
+    'reference_paths',
+    pg_typing.List(pg_typing.Object(utils.KeyPath)),
+    (
+        'Paths of referenced values, which are relative paths searched from '
+        'current node to root.'
+    ),
+)])
 class DerivedValue(symbolic.Object, pg_typing.CustomTyping):
   """Base class of value that references to other values in object tree."""
 
@@ -36,8 +39,10 @@ class DerivedValue(symbolic.Object, pg_typing.CustomTyping):
 
   def resolve(
       self, reference_path_or_paths: Optional[Union[str, List[str]]] = None
-      ) -> Union[Tuple[symbolic.Symbolic, object_utils.KeyPath],
-                 List[Tuple[symbolic.Symbolic, object_utils.KeyPath]]]:
+  ) -> Union[
+      Tuple[symbolic.Symbolic, utils.KeyPath],
+      List[Tuple[symbolic.Symbolic, utils.KeyPath]],
+  ]:
     """Resolve reference paths based on the location of this node.
 
     Args:
@@ -54,17 +59,17 @@ class DerivedValue(symbolic.Object, pg_typing.CustomTyping):
     if reference_path_or_paths is None:
       reference_paths = self.reference_paths
     elif isinstance(reference_path_or_paths, str):
-      reference_paths = [object_utils.KeyPath.parse(reference_path_or_paths)]
+      reference_paths = [utils.KeyPath.parse(reference_path_or_paths)]
       single_input = True
-    elif isinstance(reference_path_or_paths, object_utils.KeyPath):
+    elif isinstance(reference_path_or_paths, utils.KeyPath):
       reference_paths = [reference_path_or_paths]
       single_input = True
     elif isinstance(reference_path_or_paths, list):
       paths = []
       for path in reference_path_or_paths:
         if isinstance(path, str):
-          path = object_utils.KeyPath.parse(path)
-        elif not isinstance(path, object_utils.KeyPath):
+          path = utils.KeyPath.parse(path)
+        elif not isinstance(path, utils.KeyPath):
           raise ValueError('Argument \'reference_path_or_paths\' must be None, '
                            'a string, KeyPath object, a list of strings, or a '
                            'list of KeyPath objects.')
@@ -96,8 +101,7 @@ class DerivedValue(symbolic.Object, pg_typing.CustomTyping):
       # Make sure referenced value does not have referenced value.
       # NOTE(daiyip): We can support dependencies between derived values
       # in future if needed.
-      if not object_utils.traverse(
-          referenced_value, self._contains_not_derived_value):
+      if not utils.traverse(referenced_value, self._contains_not_derived_value):
         raise ValueError(
             f'Derived value (path={referenced_value.sym_path}) should not '
             f'reference derived values. '
@@ -107,15 +111,18 @@ class DerivedValue(symbolic.Object, pg_typing.CustomTyping):
     return self.derive(*referenced_values)
 
   def _contains_not_derived_value(
-      self, path: object_utils.KeyPath, value: Any) -> bool:
+      self, path: utils.KeyPath, value: Any
+  ) -> bool:
     """Returns whether a value contains derived value."""
     if isinstance(value, DerivedValue):
       return False
     elif isinstance(value, symbolic.Object):
       for k, v in value.sym_items():
-        if not object_utils.traverse(
-            v, self._contains_not_derived_value,
-            root_path=object_utils.KeyPath(k, path)):
+        if not utils.traverse(
+            v,
+            self._contains_not_derived_value,
+            root_path=utils.KeyPath(k, path),
+        ):
           return False
     return True
 
@@ -137,12 +144,13 @@ class ValueReference(DerivedValue):
 
   def custom_apply(
       self,
-      path: object_utils.KeyPath,
+      path: utils.KeyPath,
       value_spec: pg_typing.ValueSpec,
       allow_partial: bool,
-      child_transform: Optional[Callable[
-          [object_utils.KeyPath, pg_typing.Field, Any], Any]] = None
-      ) -> Tuple[bool, 'DerivedValue']:
+      child_transform: Optional[
+          Callable[[utils.KeyPath, pg_typing.Field, Any], Any]
+      ] = None,
+  ) -> Tuple[bool, 'DerivedValue']:
     """Implement pg_typing.CustomTyping interface."""
     # TODO(daiyip): perform possible static analysis on referenced paths.
     del path, value_spec, allow_partial, child_transform

pyglove/core/hyper/derived_test.py

@@ -11,13 +11,11 @@
 # 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.hyper.ValueReference."""
-
 import unittest
 
-from pyglove.core import object_utils
 from pyglove.core import symbolic
 from pyglove.core import typing as pg_typing
+from pyglove.core import utils
 from pyglove.core.hyper.derived import ValueReference
 
 
@@ -47,7 +45,7 @@ class ValueReferenceTest(unittest.TestCase):
     self.assertEqual(sd.c[0].y.resolve(), [(sd.c[0], 'c[0].x[0].z')])
     self.assertEqual(sd.c[1].y.resolve(), [(sd.c[1], 'c[1].x[0].z')])
     # Resolve references from this point.
-    self.assertEqual(sd.c[0].y.resolve(object_utils.KeyPath(0)), (sd.c, 'c[0]'))
+    self.assertEqual(sd.c[0].y.resolve(utils.KeyPath(0)), (sd.c, 'c[0]'))
     self.assertEqual(sd.c[0].y.resolve('[0]'), (sd.c, 'c[0]'))
     self.assertEqual(
         sd.c[0].y.resolve(['[0]', '[1]']), [(sd.c, 'c[0]'), (sd.c, 'c[1]')])

pyglove/core/hyper/dynamic_evaluation.py

@@ -18,9 +18,9 @@ import types
 from typing import Any, Callable, Dict, Iterator, List, Optional, Union
 
 from pyglove.core import geno
-from pyglove.core import object_utils
 from pyglove.core import symbolic
 from pyglove.core import typing as pg_typing
+from pyglove.core import utils
 from pyglove.core.hyper import base
 from pyglove.core.hyper import categorical
 from pyglove.core.hyper import custom
@@ -243,7 +243,7 @@ class DynamicEvaluationContext:
     """Registers a parameter with current context and return its first value."""
     def _add_child_decision_point(c):
       if isinstance(c, types.LambdaType):
-        s = pg_typing.get_signature(c)
+        s = pg_typing.signature(c, auto_typing=False, auto_doc=False)
         if not s.args and not s.has_wildcard_args:
           sub_context = DynamicEvaluationContext(
               where=self._where, per_thread=self._per_thread)
@@ -464,7 +464,7 @@ class DynamicEvaluationContext:
     get_current_decision = self._decision_getter
     def _apply_child(c):
       if isinstance(c, types.LambdaType):
-        s = pg_typing.get_signature(c)
+        s = pg_typing.signature(c, auto_typing=False, auto_doc=False)
         if not s.args and not s.has_wildcard_args:
           return c()
       return c
@@ -520,10 +520,10 @@ class _DynamicEvaluationStack:
   @property
   def _local_stack(self):
     """Returns thread-local stack."""
-    stack = object_utils.thread_local_get(self._TLS_KEY, None)
+    stack = utils.thread_local_get(self._TLS_KEY, None)
     if stack is None:
       stack = []
-      object_utils.thread_local_set(self._TLS_KEY, stack)
+      utils.thread_local_set(self._TLS_KEY, stack)
     return stack
 
   def push(self, context: DynamicEvaluationContext):
@@ -585,4 +585,3 @@ def trace(
   with context.collect():
     fun()
   return context
-