@SuperN1ck would be nice if you could skim through this and let me know if it makes sense!
Thing that affects you: after this is merged, factor classes will need to be decorated with @register_pytree_dataclass*.
Main change is to refactor the logic used for designating static dataclass fields. Allows us to:
Specify static fields using dataclasses.field metadata.
Unify registration logic used across the library. Factors used to do their own thing.
*I may just pull all this dataclass code into its own library and do some renaming, it seems somewhat generally useful and could help reduce boilerplate in jaxlie as well.
@SuperN1ck would be nice if you could skim through this and let me know if it makes sense!
Thing that affects you: after this is merged, factor classes will need to be decorated with
@register_pytree_dataclass*.Main change is to refactor the logic used for designating static dataclass fields. Allows us to:
dataclasses.fieldmetadata.*I may just pull all this dataclass code into its own library and do some renaming, it seems somewhat generally useful and could help reduce boilerplate in jaxlie as well.