wrappers

Author: Jammy2211

autolens/mock.py

@@ -1,6 +1,6 @@
 import jax.numpy as jnp
 
-from autofit.jax_wrapper import register_pytree_node_class
+from autoconf.jax_wrapper import register_pytree_node_class
 from autofit.mock import *  # noqa
 from autoarray.mock import *  # noqa
 from autogalaxy.mock import *  # noqa

autolens/point/solver/point_solver.py

@@ -5,7 +5,7 @@
 import autoarray as aa
 from autoarray.structures.triangles.shape import Point
 
-from autofit.jax_wrapper import register_pytree_node_class
+from autoconf.jax_wrapper import register_pytree_node_class
 from autogalaxy import OperateDeflections
 from .shape_solver import AbstractSolver
 

autolens/point/solver/shape_solver.py

@@ -7,7 +7,7 @@
 import autoarray as aa
 
 from autoarray.structures.triangles.shape import Shape
-from autofit.jax_wrapper import register_pytree_node_class
+from autoconf.jax_wrapper import register_pytree_node_class
 
 from autoarray.structures.triangles.coordinate_array import (
     CoordinateArrayTriangles,

autolens/point/solver/step.py

@@ -1,8 +1,8 @@
 import logging
 from dataclasses import dataclass
+from jax._src.tree_util import register_pytree_node_class
 
 import autoarray as aa
-from autoarray.numpy_wrapper import register_pytree_node_class
 
 logger = logging.getLogger(__name__)