fixing bugs

utils.py

@@ -4,48 +4,49 @@ import jax.numpy as jnp
 import numpy as np
 
 
-def repeat_vmap(fun, in_axes=None):
-    if in_axes is None:
-        in_axes = [0]
+def repeat_vmap(fun, 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 segurança numérica"""
+    # Garantir tamanho mínimo de 8x8
     if isinstance(patch_size, int):
-        patch_size = (max(1, patch_size), max(1, patch_size))
+        h = w = max(8, patch_size)
+    else:
+        h, w = (max(8, ps) for ps in patch_size)
 
-    offset_h, offset_w = 1 / (2 * np.array(patch_size))
-    space_h = np.linspace(-0.5 + offset_h, 0.5 - offset_h, patch_size[0])
-    space_w = np.linspace(-0.5 + offset_w, 0.5 - offset_w, patch_size[1])
+    # Espaçamento preciso entre pontos
+    y_space = np.linspace(-0.5 + 1 / (2 * h), 0.5 - 1 / (2 * h), h)
+    x_space = np.linspace(-0.5 + 1 / (2 * w), 0.5 - 1 / (2 * w), w)
 
-    grid = np.stack(np.meshgrid(space_h, space_w, indexing='ij'), axis=-1)
-    return grid[np.newaxis, ...]  # Adiciona dimensão de batch
+    # Criar grid com dimensões (1, H, W, 2)
+    grid = np.stack(np.meshgrid(y_space, x_space, indexing='ij'), axis=-1)
+    return grid[np.newaxis, ...]
 
 
 def interpolate_grid(coords, grid, order=0):
-    """Args:
-        coords: Tensor de shape (B, H, W, 2) ou (H, W, 2)
-        grid: Tensor de shape (B, H', W', C)
-        order: default 0
-    """
+    """Interpolação segura com verificação de dimensões"""
     try:
-        # Converter para array JAX e ajustar dimensões
+        # Converter para JAX array e validar formato
         coords = jnp.asarray(coords)
-        while coords.ndim < 4:
-            coords = coords[jnp.newaxis, ...]
-
-        # Verificação final de dimensões
-        if coords.shape[-1] != 2 or coords.ndim != 4:
-            raise ValueError(f"Formato inválido: {coords.shape}. Esperado (B, H, W, 2)")
+        if coords.ndim != 4 or coords.shape[-1] != 2:
+            raise ValueError(
+                f"Dimensões inválidas: {coords.shape}. Esperado (B, H, W, 2)"
+            )
 
         # Transformação de coordenadas
         coords = coords.transpose((0, 3, 1, 2))
-        coords = coords.at[:, 0].set(coords[:, 0] * grid.shape[-3] + (grid.shape[-3] - 1) / 2)
-        coords = coords.at[:, 1].set(coords[:, 1] * grid.shape[-2] + (grid.shape[-2] - 1) / 2)
-
-        # Função de interpolação vetorizada
+        coords = coords.at[:, 0].set(
+            coords[:, 0] * (grid.shape[-3] - 1) + (grid.shape[-3] - 1) / 2
+        )
+        coords = coords.at[:, 1].set(
+            coords[:, 1] * (grid.shape[-2] - 1) + (grid.shape[-2] - 1) / 2
+        )
+
+        # Interpolação vetorizada
         map_fn = jax.vmap(jax.vmap(
             partial(jax.scipy.ndimage.map_coordinates, order=order, mode='nearest'),
             in_axes=(2, None),
@@ -54,4 +55,4 @@ def interpolate_grid(coords, grid, order=0):
         return map_fn(grid, coords)
 
     except Exception as e:
-        raise RuntimeError(f"Falha na interpolação: {str(e)}") from e
+        raise RuntimeError(f"Erro de interpolação: {str(e)}") from e