Merge pull request #360 from PyAutoLabs/feature/linear-light-profile-intensity-dict-pytree

feat: stable pytree_token on LightProfileLinear for jax.jit round-trip

Author: Jammy2211

autogalaxy/profiles/light/linear/abstract.py

@@ -13,6 +13,7 @@
 intensities to be penalized by a smoothness prior.
 """
 import inspect
+import itertools
 import numpy as np
 from typing import Dict, List, Optional
 
@@ -45,6 +46,22 @@ class LightProfileLinear(LightProfile):
     profiles when performing a fit to data.
     """
 
+    _pytree_token_counter = itertools.count()
+    __exclude_identifier_fields__ = ("pytree_token",)
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.pytree_token = next(LightProfileLinear._pytree_token_counter)
+
+    def __hash__(self):
+        return self.pytree_token
+
+    def __eq__(self, other):
+        return (
+            isinstance(other, LightProfileLinear)
+            and self.pytree_token == other.pytree_token
+        )
+
     @property
     def regularization(self):
         return None