fixing bugs

utils.py

@@ -5,6 +5,14 @@ import numpy as np
 from functools import partial
 
 
+def repeat_vmap(fun, in_axes=None):
+    if in_axes is None:
+        in_axes = [0]
+    for axes in in_axes:
+        fun = jax.vmap(fun, in_axes=axes)
+    return fun
+
+
 def make_grid(patch_size: int | tuple[int, int]):
     """Gera grid de coordenadas com validação robusta"""
     if isinstance(patch_size, int):