This PR enables adding Field to instance variables of treeclass wrapped classes after initialization using .at functionality.
Example, in torch we can do the following:
Torch overrides setattr to recognize values of Parameter andModule type to register them as model parameters.
The examples show the ability to add parameters not defined in the init method; instead, it's being added after initialization.
Since PyTreeClass wrapped classes are immutable by default, setting an attribute after initialization is only possible using the .at functionality. .at creates a new instance with the updated value. However, in pytreeclass, fields (I.e. class parameters) are tied to the class, not the instance. This means no field can be added after initialization through the .at method. This PR adds this functionality, which is demonstrated in the following example:
import pytreeclass as pytc
from typing import Any
@pytc.treeclass
class Parameter:
value: Any
@pytc.treeclass
class Tree:
bias : int = 0
def add_param(self, name, param):
return setattr(self, name, param)
tree = Tree()
_, tree_with_weight_param = tree.at['add_param']('weight', Parameter(3))
print(tree)
# Tree(bias=0)
print(tree_with_weight_param)
# Tree(bias=0, weight=Parameter(value=3))
This PR enables adding
Field
to instance variables oftreeclass
wrapped classes after initialization using.at
functionality. Example, in torch we can do the following:Torch overrides
setattr
to recognize values ofParameter
andModule
type to register them as model parameters. The examples show the ability to add parameters not defined in the init method; instead, it's being added after initialization.Since
PyTreeClass
wrapped classes are immutable by default, setting an attribute after initialization is only possible using the.at
functionality..at
creates a new instance with the updated value. However, in pytreeclass, fields (I.e. class parameters) are tied to the class, not the instance. This means no field can be added after initialization through the.at
method. This PR adds this functionality, which is demonstrated in the following example: