From 32378d1cb465020eee2152c990ae9e75f8e55e22 Mon Sep 17 00:00:00 2001 From: chinganc Date: Tue, 14 Jan 2025 16:43:17 -0800 Subject: [PATCH 1/8] Update to version 0.1.3.4 Fix multiple bugs 1. Sometimes TraceGraph does not follow heap ordering. 2. When @bundle(trainable=True) is applied to a class method, different class instances will unintentionally share the same parameter. 3. When calling @bundle class method(s) multiple times of a class instance, duplicated nodes of the class instance are created, which all point to the same class instance. 4. When accessing the bundled method of a class instance, it returns a partial function instead of FunModule, which makes accessing properties of the bundled method difficult. --- opto/trace/bundle.py | 20 ++++++++++++-- opto/trace/containers.py | 6 +++-- opto/version.py | 2 +- tests/unit_tests/test_class_method.py | 39 +++++++++++++++++++++++++++ 4 files changed, 62 insertions(+), 5 deletions(-) create mode 100644 tests/unit_tests/test_class_method.py diff --git a/opto/trace/bundle.py b/opto/trace/bundle.py index 8c1115fa..1bc45858 100644 --- a/opto/trace/bundle.py +++ b/opto/trace/bundle.py @@ -503,9 +503,25 @@ def wrap(self, output: Any, inputs: Union[List[Node], Dict[str, Node]], external def is_valid_output(output): return isinstance(output, Node) or (isinstance(output, tuple) and all([isinstance(o, Node) for o in output])) - def __get__(self, obj, objtype): + + # Define __set_name__ and __get__ for FunModule to act as a descriptor. + def __get__(self, obj, db_type): + if obj is None: # class method + return self # Support instance methods. - return functools.partial(self.__call__, obj) + method_name = f'__TRACE_RESERVED_bundle_{self.name}' # NOTE we assume these are secret names not taken + obj_node_name = f'__TRACE_RESERVED_self_node' + if not hasattr(obj, obj_node_name): + setattr(obj, obj_node_name, node(obj)) + if not hasattr(obj, method_name): + funmodule = copy.deepcopy(self) + funmodule.forward = functools.partial(self.forward, getattr(obj, obj_node_name)) + setattr(obj, method_name, funmodule) + fun = getattr(obj, method_name) + assert fun is not self # self is defined in the class level + assert isinstance(fun, FunModule), f"Expected {method_name} to be a FunModule, but got {type(fun)}" + # fun = functools.partial(self.__call__, obj) + return fun def detach(self): return copy.deepcopy(self) diff --git a/opto/trace/containers.py b/opto/trace/containers.py index 63e86789..8dddc2d7 100644 --- a/opto/trace/containers.py +++ b/opto/trace/containers.py @@ -38,17 +38,19 @@ def parameters_dict(self): """ parameters = {} for name, attr in inspect.getmembers(self): + if name.startswith('__TRACE_RESERVED_'): + # These are reserved for internal use. + continue if isinstance(attr, functools.partial): # this is a class method method = attr.func.__self__ if trainable_method(method): parameters[name] = method.parameter - elif trainable_method(attr): # method attribute + if trainable_method(attr): # method attribute parameters[name] = attr.parameter elif isinstance(attr, ParameterNode): parameters[name] = attr elif isinstance(attr, ParameterContainer): parameters[name] = attr - assert all(isinstance(v, (ParameterNode, ParameterContainer)) for v in parameters.values()) return parameters # include both trainable and non-trainable parameters diff --git a/opto/version.py b/opto/version.py index 223cabec..df98922f 100644 --- a/opto/version.py +++ b/opto/version.py @@ -1 +1 @@ -__version__ = "0.1.3.3" \ No newline at end of file +__version__ = "0.1.3.4" \ No newline at end of file diff --git a/tests/unit_tests/test_class_method.py b/tests/unit_tests/test_class_method.py new file mode 100644 index 00000000..5722fb27 --- /dev/null +++ b/tests/unit_tests/test_class_method.py @@ -0,0 +1,39 @@ +from opto import trace + + +@trace.model +class Model: + + @trace.bundle(trainable=True) + def forward(self, x): + return x + 1 + + +m1 = Model() +m2 = Model() +try: + assert m1.__TRACE_RESERVED_self_node != m2.__TRACE_RESERVED_self_node +except AttributeError: + # These secrets attributes are not defined yet. They will only be defined after the bundled method is accessed. + pass + +assert len(m1.parameters()) == 1 +assert len(m2.parameters()) == 1 + +assert m1.__TRACE_RESERVED_self_node != m2.__TRACE_RESERVED_self_node # they are defined now + +# each instance has a version different from the class' version +assert m1.forward != m2.forward +assert m1.forward != Model.forward +assert m2.forward.parameter == Model.forward.parameter == m1.forward.parameter + +y1 = m1.forward(1) +y2 = m1.forward(2) + +# self is not duplicated +assert m1.__TRACE_RESERVED_self_node in y1.parents +assert m1.__TRACE_RESERVED_self_node in y2.parents +assert m1.forward.parameter in y1.parents +assert m1.forward.parameter in y2.parents +assert len(y1.parents) == 3 # since it's trainable +assert len(y2.parents) == 3 From e026fbd88065d2f650b09e3c48d42ce7d7a207c5 Mon Sep 17 00:00:00 2001 From: Ching-An Cheng Date: Thu, 16 Jan 2025 11:27:32 -0800 Subject: [PATCH 2/8] Update python-app.yml to use python 3.9 --- .github/workflows/python-app.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/python-app.yml b/.github/workflows/python-app.yml index 8074be85..bda57a97 100644 --- a/.github/workflows/python-app.yml +++ b/.github/workflows/python-app.yml @@ -19,10 +19,10 @@ jobs: steps: - uses: actions/checkout@v4 - - name: Set up Python 3.10 + - name: Set up Python 3.9 uses: actions/setup-python@v3 with: - python-version: "3.10" + python-version: "3.9" - name: Install dependencies run: | python -m pip install --upgrade pip From 8f6118053151f5ace25fdf8d6c9a886d98d283b4 Mon Sep 17 00:00:00 2001 From: chinganc Date: Sat, 18 Jan 2025 10:34:34 -0800 Subject: [PATCH 3/8] Fix bug in unit test. --- tests/unit_tests/test_class_method.py | 26 ++++++++++++++++++++------ 1 file changed, 20 insertions(+), 6 deletions(-) diff --git a/tests/unit_tests/test_class_method.py b/tests/unit_tests/test_class_method.py index 5722fb27..55a28f42 100644 --- a/tests/unit_tests/test_class_method.py +++ b/tests/unit_tests/test_class_method.py @@ -11,16 +11,25 @@ def forward(self, x): m1 = Model() m2 = Model() + +# Make sure the parameters are different try: - assert m1.__TRACE_RESERVED_self_node != m2.__TRACE_RESERVED_self_node + assert m1.__TRACE_RESERVED_self_node is not m2.__TRACE_RESERVED_self_node except AttributeError: # These secrets attributes are not defined yet. They will only be defined after the bundled method is accessed. pass +# The hidden nodes are defined now assert len(m1.parameters()) == 1 assert len(m2.parameters()) == 1 -assert m1.__TRACE_RESERVED_self_node != m2.__TRACE_RESERVED_self_node # they are defined now +# Make sure the parameters are different +assert m1.__TRACE_RESERVED_self_node is not m2.__TRACE_RESERVED_self_node # they are defined now +assert m1.parameters()[0] is not m2.parameters()[0] + +# check that the reserved node is the returned parameter +assert getattr(m1, '__TRACE_RESERVED_bundle_Model.forward').parameter is m1.parameters()[0] +assert getattr(m2, '__TRACE_RESERVED_bundle_Model.forward').parameter is m2.parameters()[0] # each instance has a version different from the class' version assert m1.forward != m2.forward @@ -30,10 +39,15 @@ def forward(self, x): y1 = m1.forward(1) y2 = m1.forward(2) +from opto.trace.utils import contain # self is not duplicated -assert m1.__TRACE_RESERVED_self_node in y1.parents -assert m1.__TRACE_RESERVED_self_node in y2.parents -assert m1.forward.parameter in y1.parents -assert m1.forward.parameter in y2.parents +assert contain(y1.parents, m1.__TRACE_RESERVED_self_node) +assert contain(y2.parents, m1.__TRACE_RESERVED_self_node) +# assert m1.__TRACE_RESERVED_self_node in y1.parents +# assert m1.__TRACE_RESERVED_self_node in y2.parents +assert contain(y1.parents, m1.forward.parameter) +assert contain(y2.parents, m1.forward.parameter) +# assert m1.forward.parameter in y1.parents +# assert m1.forward.parameter in y2.parents assert len(y1.parents) == 3 # since it's trainable assert len(y2.parents) == 3 From 1c89792d9d712c5a7e9ce18c00311e198d101d29 Mon Sep 17 00:00:00 2001 From: chinganc Date: Sat, 18 Jan 2025 10:55:23 -0800 Subject: [PATCH 4/8] Fix a wrong binding bug in __get__ --- opto/trace/bundle.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/opto/trace/bundle.py b/opto/trace/bundle.py index 1bc45858..97cf0a47 100644 --- a/opto/trace/bundle.py +++ b/opto/trace/bundle.py @@ -28,7 +28,7 @@ def bundle( ): """Wrap a function as a FunModule which returns node objects. - The input signature to the wrapped function stays the same. bundle can be used with other decorators + The input signature to the wrapped function stays the same. bundle can be used with other decorators so long as they are not named 'bundle'. Args: @@ -514,8 +514,8 @@ def __get__(self, obj, db_type): if not hasattr(obj, obj_node_name): setattr(obj, obj_node_name, node(obj)) if not hasattr(obj, method_name): - funmodule = copy.deepcopy(self) - funmodule.forward = functools.partial(self.forward, getattr(obj, obj_node_name)) + funmodule = copy.deepcopy(self) # instance specific version + funmodule.forward = functools.partial(funmodule.forward, getattr(obj, obj_node_name)) setattr(obj, method_name, funmodule) fun = getattr(obj, method_name) assert fun is not self # self is defined in the class level From 5a34e7df38d28c35eaf0faac0c90b23fa117125b Mon Sep 17 00:00:00 2001 From: windweller Date: Sat, 18 Jan 2025 11:20:19 -0800 Subject: [PATCH 5/8] add a copy test --- tests/unit_tests/test_class_method.py | 134 +++++++++++++++++--------- 1 file changed, 91 insertions(+), 43 deletions(-) diff --git a/tests/unit_tests/test_class_method.py b/tests/unit_tests/test_class_method.py index 55a28f42..9106490d 100644 --- a/tests/unit_tests/test_class_method.py +++ b/tests/unit_tests/test_class_method.py @@ -1,5 +1,5 @@ from opto import trace - +from copy import deepcopy @trace.model class Model: @@ -9,45 +9,93 @@ def forward(self, x): return x + 1 -m1 = Model() -m2 = Model() - -# Make sure the parameters are different -try: - assert m1.__TRACE_RESERVED_self_node is not m2.__TRACE_RESERVED_self_node -except AttributeError: - # These secrets attributes are not defined yet. They will only be defined after the bundled method is accessed. - pass - -# The hidden nodes are defined now -assert len(m1.parameters()) == 1 -assert len(m2.parameters()) == 1 - -# Make sure the parameters are different -assert m1.__TRACE_RESERVED_self_node is not m2.__TRACE_RESERVED_self_node # they are defined now -assert m1.parameters()[0] is not m2.parameters()[0] - -# check that the reserved node is the returned parameter -assert getattr(m1, '__TRACE_RESERVED_bundle_Model.forward').parameter is m1.parameters()[0] -assert getattr(m2, '__TRACE_RESERVED_bundle_Model.forward').parameter is m2.parameters()[0] - -# each instance has a version different from the class' version -assert m1.forward != m2.forward -assert m1.forward != Model.forward -assert m2.forward.parameter == Model.forward.parameter == m1.forward.parameter - -y1 = m1.forward(1) -y2 = m1.forward(2) - -from opto.trace.utils import contain -# self is not duplicated -assert contain(y1.parents, m1.__TRACE_RESERVED_self_node) -assert contain(y2.parents, m1.__TRACE_RESERVED_self_node) -# assert m1.__TRACE_RESERVED_self_node in y1.parents -# assert m1.__TRACE_RESERVED_self_node in y2.parents -assert contain(y1.parents, m1.forward.parameter) -assert contain(y2.parents, m1.forward.parameter) -# assert m1.forward.parameter in y1.parents -# assert m1.forward.parameter in y2.parents -assert len(y1.parents) == 3 # since it's trainable -assert len(y2.parents) == 3 +def test_case_two_models(): + m1 = Model() + m2 = Model() + + # Make sure the parameters are different + try: + assert m1.__TRACE_RESERVED_self_node is not m2.__TRACE_RESERVED_self_node + except AttributeError: + # These secrets attributes are not defined yet. They will only be defined after the bundled method is accessed. + pass + + # The hidden nodes are defined now + assert len(m1.parameters()) == 1 + assert len(m2.parameters()) == 1 + + # Make sure the parameters are different + assert m1.__TRACE_RESERVED_self_node is not m2.__TRACE_RESERVED_self_node # they are defined now + assert m1.parameters()[0] is not m2.parameters()[0] + + # check that the reserved node is the returned parameter + assert getattr(m1, '__TRACE_RESERVED_bundle_Model.forward').parameter is m1.parameters()[0] + assert getattr(m2, '__TRACE_RESERVED_bundle_Model.forward').parameter is m2.parameters()[0] + + # each instance has a version different from the class' version + assert m1.forward != m2.forward + assert m1.forward != Model.forward + assert m2.forward.parameter == Model.forward.parameter == m1.forward.parameter + + y1 = m1.forward(1) + y2 = m1.forward(2) + + from opto.trace.utils import contain + # self is not duplicated + assert contain(y1.parents, m1.__TRACE_RESERVED_self_node) + assert contain(y2.parents, m1.__TRACE_RESERVED_self_node) + # assert m1.__TRACE_RESERVED_self_node in y1.parents + # assert m1.__TRACE_RESERVED_self_node in y2.parents + assert contain(y1.parents, m1.forward.parameter) + assert contain(y2.parents, m1.forward.parameter) + # assert m1.forward.parameter in y1.parents + # assert m1.forward.parameter in y2.parents + assert len(y1.parents) == 3 # since it's trainable + assert len(y2.parents) == 3 + +def test_case_model_copy(): + m1 = Model() + m2 = deepcopy(m1) + + # Make sure the parameters are different + try: + assert m1.__TRACE_RESERVED_self_node is not m2.__TRACE_RESERVED_self_node + except AttributeError: + # These secrets attributes are not defined yet. They will only be defined after the bundled method is accessed. + pass + + # The hidden nodes are defined now + assert len(m1.parameters()) == 1 + assert len(m2.parameters()) == 1 + + # Make sure the parameters are different + assert m1.__TRACE_RESERVED_self_node is not m2.__TRACE_RESERVED_self_node # they are defined now + assert m1.parameters()[0] is not m2.parameters()[0] + + # check that the reserved node is the returned parameter + assert getattr(m1, '__TRACE_RESERVED_bundle_Model.forward').parameter is m1.parameters()[0] + assert getattr(m2, '__TRACE_RESERVED_bundle_Model.forward').parameter is m2.parameters()[0] + + # each instance has a version different from the class' version + assert m1.forward is not m2.forward + assert m1.forward is not Model.forward + assert m2.forward.parameter == Model.forward.parameter == m1.forward.parameter + + y1 = m1.forward(1) + y2 = m1.forward(2) + + from opto.trace.utils import contain + # self is not duplicated + assert contain(y1.parents, m1.__TRACE_RESERVED_self_node) + assert contain(y2.parents, m1.__TRACE_RESERVED_self_node) + # assert m1.__TRACE_RESERVED_self_node in y1.parents + # assert m1.__TRACE_RESERVED_self_node in y2.parents + assert contain(y1.parents, m1.forward.parameter) + assert contain(y2.parents, m1.forward.parameter) + # assert m1.forward.parameter in y1.parents + # assert m1.forward.parameter in y2.parents + assert len(y1.parents) == 3 # since it's trainable + assert len(y2.parents) == 3 + +test_case_two_models() +test_case_model_copy() \ No newline at end of file From 9e2db279b1b4bfb17280bc7c5201d8b9b32672f4 Mon Sep 17 00:00:00 2001 From: windweller Date: Sat, 18 Jan 2025 11:37:20 -0800 Subject: [PATCH 6/8] push a deepcopy override --- opto/trace/modules.py | 12 +++++++++++- tests/unit_tests/test_class_method.py | 10 +++++++--- 2 files changed, 18 insertions(+), 4 deletions(-) diff --git a/opto/trace/modules.py b/opto/trace/modules.py index ae7a8913..36b0d912 100644 --- a/opto/trace/modules.py +++ b/opto/trace/modules.py @@ -61,4 +61,14 @@ def _set(self, new_parameters): parameters_dict[k]._set(v) else: # if the parameter does not exist assert k not in self.__dict__ - setattr(self, k, v) \ No newline at end of file + setattr(self, k, v) + + def __deepcopy__(self, memo): + """ Custom deepcopy behavior for Module. """ + cls = self.__class__ + result = cls.__new__(cls) + memo[id(self)] = result + for k, v in self.__dict__.items(): + if '__TRACE_RESERVED_' not in k: + setattr(result, k, copy.deepcopy(v, memo)) + return result \ No newline at end of file diff --git a/tests/unit_tests/test_class_method.py b/tests/unit_tests/test_class_method.py index 9106490d..33748f10 100644 --- a/tests/unit_tests/test_class_method.py +++ b/tests/unit_tests/test_class_method.py @@ -82,20 +82,24 @@ def test_case_model_copy(): assert m2.forward.parameter == Model.forward.parameter == m1.forward.parameter y1 = m1.forward(1) - y2 = m1.forward(2) + y2 = m2.forward(2) from opto.trace.utils import contain # self is not duplicated assert contain(y1.parents, m1.__TRACE_RESERVED_self_node) - assert contain(y2.parents, m1.__TRACE_RESERVED_self_node) + assert contain(y2.parents, m2.__TRACE_RESERVED_self_node) # assert m1.__TRACE_RESERVED_self_node in y1.parents # assert m1.__TRACE_RESERVED_self_node in y2.parents assert contain(y1.parents, m1.forward.parameter) - assert contain(y2.parents, m1.forward.parameter) + assert contain(y2.parents, m2.forward.parameter) + # assert m1.forward.parameter in y1.parents # assert m1.forward.parameter in y2.parents assert len(y1.parents) == 3 # since it's trainable assert len(y2.parents) == 3 +def printout_deecopy_modules(): + pass + test_case_two_models() test_case_model_copy() \ No newline at end of file From 85a4528b0b518ca76c0c9b74090ddcdbaa176b6b Mon Sep 17 00:00:00 2001 From: windweller Date: Sat, 18 Jan 2025 11:45:48 -0800 Subject: [PATCH 7/8] revert adding deepcopy override. Add test for nested deepcopy --- opto/trace/modules.py | 12 +------ tests/unit_tests/test_class_method.py | 52 ++++++++++++++++++++++++--- 2 files changed, 49 insertions(+), 15 deletions(-) diff --git a/opto/trace/modules.py b/opto/trace/modules.py index 36b0d912..ae7a8913 100644 --- a/opto/trace/modules.py +++ b/opto/trace/modules.py @@ -61,14 +61,4 @@ def _set(self, new_parameters): parameters_dict[k]._set(v) else: # if the parameter does not exist assert k not in self.__dict__ - setattr(self, k, v) - - def __deepcopy__(self, memo): - """ Custom deepcopy behavior for Module. """ - cls = self.__class__ - result = cls.__new__(cls) - memo[id(self)] = result - for k, v in self.__dict__.items(): - if '__TRACE_RESERVED_' not in k: - setattr(result, k, copy.deepcopy(v, memo)) - return result \ No newline at end of file + setattr(self, k, v) \ No newline at end of file diff --git a/tests/unit_tests/test_class_method.py b/tests/unit_tests/test_class_method.py index 33748f10..2bf62600 100644 --- a/tests/unit_tests/test_class_method.py +++ b/tests/unit_tests/test_class_method.py @@ -1,5 +1,5 @@ from opto import trace -from copy import deepcopy +from copy import deepcopy, copy @trace.model class Model: @@ -98,8 +98,52 @@ def test_case_model_copy(): assert len(y1.parents) == 3 # since it's trainable assert len(y2.parents) == 3 -def printout_deecopy_modules(): - pass +def test_case_model_nested_copy(): + m1 = Model() + m3 = deepcopy(m1) + m2 = deepcopy(m3) + + # Make sure the parameters are different + try: + assert m1.__TRACE_RESERVED_self_node is not m2.__TRACE_RESERVED_self_node + except AttributeError: + # These secrets attributes are not defined yet. They will only be defined after the bundled method is accessed. + pass + + # The hidden nodes are defined now + assert len(m1.parameters()) == 1 + assert len(m2.parameters()) == 1 + + # Make sure the parameters are different + assert m1.__TRACE_RESERVED_self_node is not m2.__TRACE_RESERVED_self_node # they are defined now + assert m1.parameters()[0] is not m2.parameters()[0] + + # check that the reserved node is the returned parameter + assert getattr(m1, '__TRACE_RESERVED_bundle_Model.forward').parameter is m1.parameters()[0] + assert getattr(m2, '__TRACE_RESERVED_bundle_Model.forward').parameter is m2.parameters()[0] + + # each instance has a version different from the class' version + assert m1.forward is not m2.forward + assert m1.forward is not Model.forward + assert m2.forward.parameter == Model.forward.parameter == m1.forward.parameter + + y1 = m1.forward(1) + y2 = m2.forward(2) + + from opto.trace.utils import contain + # self is not duplicated + assert contain(y1.parents, m1.__TRACE_RESERVED_self_node) + assert contain(y2.parents, m2.__TRACE_RESERVED_self_node) + # assert m1.__TRACE_RESERVED_self_node in y1.parents + # assert m1.__TRACE_RESERVED_self_node in y2.parents + assert contain(y1.parents, m1.forward.parameter) + assert contain(y2.parents, m2.forward.parameter) + + # assert m1.forward.parameter in y1.parents + # assert m1.forward.parameter in y2.parents + assert len(y1.parents) == 3 # since it's trainable + assert len(y2.parents) == 3 test_case_two_models() -test_case_model_copy() \ No newline at end of file +test_case_model_copy() +test_case_model_nested_copy() \ No newline at end of file From ab49d1a05384c83311369863bece9ca7af5fbeca Mon Sep 17 00:00:00 2001 From: chinganc Date: Sun, 19 Jan 2025 13:41:59 -0800 Subject: [PATCH 8/8] Fix the bug that __deepcopy__ does not register new node. --- opto/trace/nodes.py | 1 + 1 file changed, 1 insertion(+) diff --git a/opto/trace/nodes.py b/opto/trace/nodes.py index 4e522936..5eda2872 100644 --- a/opto/trace/nodes.py +++ b/opto/trace/nodes.py @@ -466,6 +466,7 @@ def __deepcopy__(self, memo): setattr(result, k, defaultdict(list)) else: setattr(result, k, copy.deepcopy(v, memo)) + GRAPH.register(result) return result def lt(self, other):