cf69b916af
remove "return state" for functions which will be executed in vmap; recover randkey as args in mutation methods
41 lines
1.2 KiB
Python
41 lines
1.2 KiB
Python
from jax.tree_util import register_pytree_node_class
|
|
|
|
|
|
@register_pytree_node_class
|
|
class State:
|
|
def __init__(self, **kwargs):
|
|
self.__dict__["state_dict"] = kwargs
|
|
|
|
def registered_keys(self):
|
|
return self.state_dict.keys()
|
|
|
|
def register(self, **kwargs):
|
|
for key in kwargs:
|
|
if key in self.registered_keys():
|
|
raise ValueError(f"Key {key} already exists in state")
|
|
return State(**{**self.state_dict, **kwargs})
|
|
|
|
def update(self, **kwargs):
|
|
for key in kwargs:
|
|
if key not in self.registered_keys():
|
|
raise ValueError(f"Key {key} does not exist in state")
|
|
return State(**{**self.state_dict, **kwargs})
|
|
|
|
def __getattr__(self, name):
|
|
return self.state_dict[name]
|
|
|
|
def __setattr__(self, name, value):
|
|
raise AttributeError("State is immutable")
|
|
|
|
def __repr__(self):
|
|
return f"State ({self.state_dict})"
|
|
|
|
def tree_flatten(self):
|
|
children = list(self.state_dict.values())
|
|
aux_data = list(self.state_dict.keys())
|
|
return children, aux_data
|
|
|
|
@classmethod
|
|
def tree_unflatten(cls, aux_data, children):
|
|
return cls(**dict(zip(aux_data, children)))
|