| 660 | |
| 661 | @functools.wraps(original_init) |
| 662 | def init(self, *args, **kwargs): |
| 663 | if not isinstance(self, ConfigMixin): |
| 664 | raise RuntimeError( |
| 665 | f"`@register_for_config` was applied to {self.__class__.__name__} init method, but this class does " |
| 666 | "not inherit from `ConfigMixin`." |
| 667 | ) |
| 668 | |
| 669 | # Ignore private kwargs in the init. Retrieve all passed attributes |
| 670 | init_kwargs = dict(kwargs.items()) |
| 671 | |
| 672 | # Retrieve default values |
| 673 | fields = dataclasses.fields(self) |
| 674 | default_kwargs = {} |
| 675 | for field in fields: |
| 676 | # ignore flax specific attributes |
| 677 | if field.name in self._flax_internal_args: |
| 678 | continue |
| 679 | if type(field.default) == dataclasses._MISSING_TYPE: |
| 680 | default_kwargs[field.name] = None |
| 681 | else: |
| 682 | default_kwargs[field.name] = getattr(self, field.name) |
| 683 | |
| 684 | # Make sure init_kwargs override default kwargs |
| 685 | new_kwargs = {**default_kwargs, **init_kwargs} |
| 686 | # dtype should be part of `init_kwargs`, but not `new_kwargs` |
| 687 | if "dtype" in new_kwargs: |
| 688 | new_kwargs.pop("dtype") |
| 689 | |
| 690 | # Get positional arguments aligned with kwargs |
| 691 | for i, arg in enumerate(args): |
| 692 | name = fields[i].name |
| 693 | new_kwargs[name] = arg |
| 694 | |
| 695 | # Take note of the parameters that were not present in the loaded config |
| 696 | if len(set(new_kwargs.keys()) - set(init_kwargs)) > 0: |
| 697 | new_kwargs["_use_default_values"] = list(set(new_kwargs.keys()) - set(init_kwargs)) |
| 698 | |
| 699 | getattr(self, "register_to_config")(**new_kwargs) |
| 700 | original_init(self, *args, **kwargs) |
| 701 | |
| 702 | cls.__init__ = init |
| 703 | return cls |