Merge pull request #1037 from rhayes777/feature/jax_pytree

fix setting attribute rather than calling in constructor. could cause issues later

autofit/mapper/prior_model/prior_model.py

@@ -245,9 +245,17 @@ def instance_flatten(self, instance):
         """
         Flatten an instance of this model as a PyTree.
         """
+        attribute_names = [
+            name
+            for name in self.direct_argument_names
+            if hasattr(instance, name) and name not in self.constructor_argument_names
+        ]
         return (
-            [getattr(instance, name) for name in self.direct_argument_names],
-            None,
+            (
+                [getattr(instance, name) for name in self.constructor_argument_names],
+                [getattr(instance, name) for name in attribute_names],
+            ),
+            (attribute_names,),
         )
 
     def instance_unflatten(self, aux_data, children):
@@ -263,7 +271,12 @@ def instance_unflatten(self, aux_data, children):
         -------
         An instance of this model.
         """
-        return self.cls(**dict(zip(self.direct_argument_names, children)))
+        constructor_arguments, other_arguments = children
+        attribute_names = aux_data[0]
+        instance = self.cls(*constructor_arguments)
+        for name, value in zip(attribute_names, other_arguments):
+            setattr(instance, name, value)
+        return instance
 
     def tree_flatten(self):
         """