[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/API/server/main_server_pro.py

@@ -58,10 +58,15 @@ def compute_gempy_model(gempy_input: GemPyInput) -> Response:
         logger=logger
     )
 
+    # Options contain mutable nested configuration, so each request needs its
+    # own copy. The request schema exposes the desired octree depth on the grid.
+    options = default_interpolation_options.model_copy(deep=True)
+    options.evaluation_options.number_octree_levels = gempy_input.interpolation_input.grid.octree_levels
+
     # Compute model
     solutions = _compute_model(
         interpolation_input=interpolation_input,
-        options=default_interpolation_options,
+        options=options,
         structure=input_data_descriptor
     )
     logger.info("Finished computing model")
@@ -107,4 +112,4 @@ def _compute_model(
             v_just_points=interpolation_input.surface_points.sp_coords
         )
 
-    return solutions
+    return solutions

gempy_engine/core/backend_tensor.py

@@ -195,8 +195,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.t.array(list(range(tensors_structure.n_surfaces)), dtype=rest_nugget.dtype),
         number_repetitions,
         0,
     )

gempy_engine/modules/evaluator/symbolic_evaluator.py

@@ -104,6 +104,21 @@ def _build_block_sparse_ranges(M_sizes: list[int], N_sizes: list[int]):
     return numpy_ranges
 
 
+def _validate_stacked_dimensions(eval_kernel, weights, M_sizes: list[int], N_sizes: list[int]) -> None:
+    """Fail before PyKeOps when block ranges do not describe the lazy kernel."""
+    expected_i = sum(N_sizes)
+    expected_j = sum(M_sizes)
+    kernel_shape = tuple(eval_kernel.shape)
+    weights_size = weights.shape[0]
+
+    if kernel_shape[0] != expected_i or kernel_shape[1] != expected_j or weights_size != expected_i:
+        raise ValueError(
+            "Inconsistent stacked PyKeOps dimensions: "
+            f"kernel={kernel_shape}, weights={weights_size}, "
+            f"range dimensions=({expected_i}, {expected_j})."
+        )
+
+
 def symbolic_evaluator_optimized_stacked(
         eval_inputs: list[EvaluatorInput],
         weights_list: list[np.ndarray],
@@ -194,6 +209,7 @@ def _run_prep(args):
     # Concatenate weights
     all_weights = BackendTensor.t.concatenate(weights_list, axis=0)
     all_weights = BackendTensor.t.tile(all_weights, tile_factor)
+    _validate_stacked_dimensions(eval_kernel, all_weights, M_sizes, N_sizes)
 
     if BackendTensor.engine_backend == gempy_engine.config.AvailableBackends.numpy:
         from pykeops.numpy import LazyTensor
@@ -234,13 +250,11 @@ def _run_prep(args):
         except TypeError:
             raise ValueError("Failed to compute symbolic evaluation with PyKeOps. Ensure that all_weights and eval_kernel are compatible for lazy tensor operations.")
 
-    # For torch
-    # all_results_concat = all_results_concat.to("cpu")
-    all_results_split = BackendTensor.t.split(all_results_concat, M_sizes)
-
-    # For numpy
-    # split_indices = np.cumsum(M_sizes)[:-1]
-    # all_results_split = np.split(all_results_concat, split_indices)
+    if BackendTensor.engine_backend == gempy_engine.config.AvailableBackends.numpy:
+        split_indices = np.cumsum(M_sizes)[:-1]
+        all_results_split = np.split(all_results_concat, split_indices)
+    else:
+        all_results_split = BackendTensor.t.split(all_results_concat, M_sizes)
 
     original_n_fields = len(eval_inputs)
     scalar_fields = all_results_split[:original_n_fields]

tests/test_common/test_api/test_concurrent_options.py

@@ -0,0 +1,60 @@
+from concurrent.futures import ThreadPoolExecutor
+from threading import Barrier
+
+import pytest
+
+from gempy_engine.API.model import model_api
+from gempy_engine.core.data import InterpolationOptions
+from gempy_engine.core.data.options.temp_interpolation_values import TempInterpolationValues
+from gempy_engine.modules.evaluator.symbolic_evaluator import _validate_stacked_dimensions
+
+
+def test_computations_do_not_share_volatile_option_state(monkeypatch):
+    options = InterpolationOptions.from_args(
+        range=1.0,
+        c_o=1.0,
+        number_octree_levels=2,
+        mesh_extraction=False,
+    )
+    barrier = Barrier(2)
+
+    monkeypatch.setattr(model_api.WeightCache, "initialize_cache_dir", lambda: None)
+    monkeypatch.setattr(model_api.WeightCache, "clear_cache", lambda: None)
+    monkeypatch.setattr(model_api.BackendTensor, "clear_gpu_memory", lambda: None)
+    monkeypatch.setattr(model_api, "_check_input_validity", lambda *_: None)
+
+    def observe_options(interpolation_input, options, data_descriptor):
+        level = interpolation_input
+        options.temp_interpolation_values.current_octree_level = level
+        barrier.wait()
+        return options.temp_interpolation_values.current_octree_level
+
+    monkeypatch.setattr(model_api, "interpolate_n_octree_levels", observe_options)
+
+    class FakeSolutions:
+        def __init__(self, octrees_output, **_):
+            self.observed_level = octrees_output
+
+    monkeypatch.setattr(model_api, "Solutions", FakeSolutions)
+
+    with ThreadPoolExecutor(max_workers=2) as executor:
+        futures = [
+            executor.submit(model_api.compute_model, level, options, None)
+            for level in (1, 3)
+        ]
+
+    assert [future.result().observed_level for future in futures] == [1, 3]
+    assert options.temp_interpolation_values == TempInterpolationValues()
+
+
+def test_stacked_dimension_validation_reports_range_mismatch():
+    class FakeKernel:
+        shape = (4, 8, 1)
+
+    class FakeWeights:
+        shape = (4,)
+
+    with pytest.raises(ValueError, match="range dimensions=\\(4, 16\\)") as error:
+        _validate_stacked_dimensions(FakeKernel(), FakeWeights(), M_sizes=[8, 8], N_sizes=[2, 2])
+
+    assert "kernel=(4, 8, 1)" in str(error.value)

tests/test_common/test_api/test_stack_options_override.py

@@ -86,8 +86,10 @@ def test_stack_options_override_serial(override_setup):
 
 
 @pytest.mark.skipif(not PYKEOPS_AVAILABLE, reason="pykeops not installed")
-def test_stack_options_override_flat(override_setup):
+@pytest.mark.parametrize("number_octree_levels", [1, 2])
+def test_stack_options_override_flat(override_setup, number_octree_levels):
     ii, global_options, ts = override_setup
+    global_options.evaluation_options.number_octree_levels = number_octree_levels
 
     custom_options = InterpolationOptions.from_args(
         range=5.0,

tests/test_common/test_modules/test_kernel_constructor/test_nugget_propagation.py

@@ -125,6 +125,36 @@ def test_surface_preprocessing_preserves_nugget_components_and_surface_ids(simpl
     )
 
 
+def test_surface_preprocessing_uses_backend_device_when_torch_default_differs(simple_model_2):
+    if BackendTensor.engine_backend is not AvailableBackends.PYTORCH or not BackendTensor.use_gpu:
+        pytest.skip("PyTorch GPU-only device regression test")
+    import torch
+
+    previous_default_device = torch.get_default_device()
+    try:
+        torch.set_default_device("cpu")
+        surface_points, _, _, descriptor = deepcopy(simple_model_2)
+
+        internal = surface_points_preprocess(surface_points, descriptor.tensors_structure)
+
+        assert internal.surface_ids.device.type == BackendTensor.device.type
+    finally:
+        torch.set_default_device(previous_default_device)
+
+
+def test_pytorch_repeat_moves_existing_tensors_to_backend_device():
+    if BackendTensor.engine_backend is not AvailableBackends.PYTORCH or not BackendTensor.use_gpu:
+        pytest.skip("PyTorch GPU-only device regression test")
+    import torch
+
+    values = torch.tensor([0.0], device="cpu")
+    repeats = torch.tensor([1], device=BackendTensor.device)
+
+    result = BackendTensor.t.repeat(values, repeats, 0)
+
+    assert result.device.type == BackendTensor.device.type
+
+
 def test_all_nuggets_receive_pytorch_gradients(simple_model_2):
     if BackendTensor.engine_backend is not AvailableBackends.PYTORCH:
         pytest.skip("PyTorch-only autograd test")