build: raise astropy cap, update JAX Partial imports for 0.5+ compat

autogalaxy/operate/lens_calc.py

@@ -983,7 +983,10 @@ def _make_eigen_fn(self, kind: str, pixel_scales=(0.05, 0.05)):
         """
         import jax
         import jax.numpy as jnp
-        from jax.tree_util import Partial
+        try:
+            from jax.tree_util import Partial
+        except ImportError:
+            from functools import partial as Partial
 
         # Capture as local names so the closure holds no `self` reference.
         # ZeroSolver.zero_contour_finder is jit-compiled with `f` as a

autogalaxy/profiles/mass/total/jax_utils.py

@@ -32,7 +32,7 @@ def omega(eiphi, slope, factor, n_terms=20, xp=np):
         be sufficient most of the time)
     """
 
-    from jax.tree_util import Partial as partial
+    from functools import partial
     import jax
 
     scan = jax.jit(jax.lax.scan, static_argnames=("length", "reverse", "unroll"))

pyproject.toml

@@ -25,7 +25,7 @@ dependencies = [
     "autofit",
     "autoarray",
     "colossus==1.3.1",
-    "astropy>=5.0,<=6.1.2",
+    "astropy>=5.0,<=7.2.0",
     "nautilus-sampler==1.0.5"
 ]