undo changes

Author: Jammy2211

autolens/mock.py

@@ -13,7 +13,6 @@
 from autolens.point.mock.mock_solver import MockPointSolver  # noqa
 
 
-@register_pytree_node_class
 class NullTracer(Tracer):
     def __init__(self):
         super().__init__([])
@@ -28,16 +27,3 @@ def magnification_2d_via_hessian_from(
         self, grid, buffer: float = 0.01, deflections_func=None
     ) -> aa.ArrayIrregular:
         return aa.ArrayIrregular(values=jnp.ones(grid.shape[0]))
-
-    def tree_flatten(self):
-        """
-        Flatten this model as a PyTree.
-        """
-        return (), None
-
-    @classmethod
-    def tree_unflatten(cls, aux_data, children):
-        """
-        Unflatten a PyTree into a model.
-        """
-        return cls()

autolens/point/solver/point_solver.py

@@ -12,8 +12,6 @@
 
 logger = logging.getLogger(__name__)
 
-
-@register_pytree_node_class
 class PointSolver(AbstractSolver):
 
     def solve(

autolens/point/solver/shape_solver.py

@@ -329,25 +329,6 @@ def steps(
 
             initial_triangles = up_sampled
 
-    def tree_flatten(self):
-        return (), (
-            self.scale,
-            self.pixel_scale_precision,
-            self.magnification_threshold,
-            self.initial_triangles,
-        )
-
-    @classmethod
-    def tree_unflatten(cls, aux_data, children):
-        return cls(
-            scale=aux_data[0],
-            pixel_scale_precision=aux_data[1],
-            magnification_threshold=aux_data[2],
-            initial_triangles=aux_data[3],
-        )
-
-
-@register_pytree_node_class
 class ShapeSolver(AbstractSolver):
     def find_magnification(
         self,

autolens/point/solver/step.py

@@ -2,12 +2,10 @@
 from dataclasses import dataclass
 
 import autoarray as aa
-from autoarray.numpy_wrapper import register_pytree_node_class
 
 logger = logging.getLogger(__name__)
 
 
-@register_pytree_node_class
 @dataclass
 class Step:
     """
@@ -33,16 +31,3 @@ class Step:
     neighbourhood: aa.AbstractTriangles
     up_sampled: aa.AbstractTriangles
     plane_triangles: aa.AbstractTriangles
-
-    def tree_flatten(self):
-        return (
-            self.number,
-            self.initial_triangles,
-            self.filtered_triangles,
-            self.neighbourhood,
-            self.up_sampled,
-        ), ()
-
-    @classmethod
-    def tree_unflatten(cls, _, values):
-        return cls(*values)