[FIX] Ensure tensors consistently move to backend device and add tests for GPU-specific behavior

gempy_engine/API/model/model_api.py

@@ -18,6 +18,7 @@
 from ...core.data.solutions import Solutions
 from ...core.utils import gempy_profiler_decorator
 from ...core.exceptions import GemPyEngineInputError
+from ...core.data.options.temp_interpolation_values import TempInterpolationValues
 from ...modules.geophysics.fw_gravity import compute_gravity
 from ...modules.geophysics.fw_magnetic import compute_tmi
 from ...modules.weights_cache.weights_cache_interface import WeightCache
@@ -26,6 +27,10 @@
 @gempy_profiler_decorator
 def compute_model(interpolation_input: InterpolationInput, options: InterpolationOptions,
                   data_descriptor: InputDataDescriptor, *, geophysics_input: Optional[GeophysicsInput] = None) -> Solutions:
+    # Octree progress and cache timestamps belong to this computation. Keeping
+    # them on a shared options object lets concurrent requests change each
+    # other's active octree level.
+    options = options.model_copy(update={"temp_interpolation_values": TempInterpolationValues()})
     try:
         WeightCache.initialize_cache_dir()
         options.temp_interpolation_values.start_computation_ts = int(time.time())

gempy_engine/core/backend_tensor.py

@@ -185,6 +185,14 @@ def describe_conf(cls):
         print(f"\n Using gpu: {cls.use_gpu}. \n")
         print(f"\n Using pykeops: {cls.pykeops_enabled}. \n")
 
+    @classmethod
+    def arange(cls, stop, *, dtype=None):
+        if cls.engine_backend == AvailableBackends.PYTORCH:
+            if isinstance(dtype, str):
+                dtype = getattr(torch, dtype)
+            return torch.arange(stop, dtype=dtype, device=cls.device)
+        return cls.t.arange(stop, dtype=dtype)
+
     @classmethod
     def _wrap_pytorch_functions(cls):
         import torch
@@ -205,8 +213,12 @@ def _sum(tensor, axis=None, dtype=None, keepdims=False):
         def _repeat(tensor, n_repeats, axis=None):
             if not isinstance(tensor, torch.Tensor):
                 tensor = torch.as_tensor(tensor, device=cls.device)
+            elif tensor.device != cls.device:
+                tensor = tensor.to(cls.device)
             if not isinstance(n_repeats, torch.Tensor):
                 n_repeats = torch.as_tensor(n_repeats, device=cls.device)
+            elif n_repeats.device != cls.device:
+                n_repeats = n_repeats.to(cls.device)
             return _true_torch_repeat_interleave(tensor, n_repeats, dim=axis)
 
         def _array(array_like, dtype=None):

gempy_engine/modules/data_preprocess/_input_preparation.py

@@ -38,7 +38,7 @@ def surface_points_preprocess(sp_input: SurfacePoints, tensors_structure: Tensor
     ref_points_repeated = b.t.repeat(ref_points, number_repetitions, 0)  # ref_points shape: (1, 3)
     ref_nugget_repeated = b.t.repeat(ref_nugget, number_repetitions, 0)
     surface_ids = b.t.repeat(
-        b.t.arange(tensors_structure.n_surfaces, dtype=rest_nugget.dtype),
+        b.arange(tensors_structure.n_surfaces, dtype=rest_nugget.dtype),
         number_repetitions,
         0,
     )