vectorized jax?

Author: Jammy2211

autoarray/structures/mesh/triangulation_2d.py

@@ -26,23 +26,32 @@ def scipy_delaunay_padded(points_np, max_simplices):
     return pts, padded
 
 
-def jax_delaunay(points):
+def delaunay_no_vmap(points):
 
     import jax
     import jax.numpy as jnp
 
+    # prevent batching
+    points = jax.lax.stop_gradient(points)
+
     N = points.shape[0]
     max_simplices = 2 * N
 
-    pts_shape = jax.ShapeDtypeStruct((N, 2), points.dtype)
-    simp_shape = jax.ShapeDtypeStruct((max_simplices, 3), jnp.int32)
+    pts_spec  = jax.ShapeDtypeStruct((N, 2), points.dtype)
+    simp_spec = jax.ShapeDtypeStruct((max_simplices, 3), jnp.int32)
+
+    def _cb(pts):
+        return scipy_delaunay_padded(pts, max_simplices)
 
-    return jax.pure_callback(
-        lambda pts: scipy_delaunay_padded(pts, max_simplices),
-        (pts_shape, simp_shape),
+    pts_out, simplices_out = jax.pure_callback(
+        _cb,
+        (pts_spec, simp_spec),
         points,
+        vectorized=False,   # ← VERY IMPORTANT (JAX 0.4.33+)
     )
 
+    return pts_out, simplices_out
+
 
 def find_simplex_from(query_points, points, simplices):
     """
@@ -146,7 +155,7 @@ def circumcenters_from(points, simplices, xp=np):
     return centers
 
 
-MAX_DEG_JAX = 32
+MAX_DEG_JAX = 64
 
 def voronoi_areas_via_delaunay_from(points, simplices, xp=np):
     """
@@ -406,7 +415,7 @@ def delaunay(self) -> "scipy.spatial.Delaunay":
 
             import jax.numpy as jnp
 
-            points, simplices = jax_delaunay(mesh_grid)
+            points, simplices = delaunay_no_vmap(mesh_grid)
             vertex_neighbor_vertices = None
 
         else: