[ENH] Add gradient computation control and improve memory management (#31)

### TL;DR

Added memory optimization features to improve performance when evaluating large models.

### What changed?

- Added a `grads` parameter to the backend tensor configuration to control PyTorch gradient computation
- Implemented proper context management for PyTorch's `no_grad()` mode
- Added garbage collection calls during chunked evaluation to prevent memory buildup
- Optimized memory usage in the evaluation kernel by immediately deleting tensors after use
- Improved error handling in the evaluation function

### How to test?

1. Test with large models that previously caused memory issues
2. Compare memory usage before and after these changes
3. Verify that model evaluation still produces correct results
4. Test with both gradient computation enabled and disabled

### Why make this change?

These optimizations address memory leaks and excessive memory usage during model evaluation, particularly for large models. By properly managing PyTorch's gradient computation and implementing strategic garbage collection, we can significantly reduce memory footprint without sacrificing performance. The immediate cleanup of tensors after use prevents memory buildup during evaluation of large datasets.

Author: Leguark

gempy_engine/core/backend_tensor.py

@@ -46,11 +46,14 @@ def get_backend_string(cls) -> str:
                 return "CPU"
 
     @classmethod
-    def change_backend_gempy(cls, engine_backend: AvailableBackends, use_gpu: bool = False, dtype: Optional[str] = None):
-        cls._change_backend(engine_backend, use_pykeops=PYKEOPS, use_gpu=use_gpu, dtype=dtype)
+    def change_backend_gempy(cls, engine_backend: AvailableBackends, use_gpu: bool = False,
+                             dtype: Optional[str] = None, grads:bool = False):
+        cls._change_backend(engine_backend, use_pykeops=PYKEOPS, use_gpu=use_gpu, dtype=dtype,
+                            grads=grads)
 
     @classmethod
-    def _change_backend(cls, engine_backend: AvailableBackends, use_pykeops: bool = False, use_gpu: bool = True, dtype: Optional[str] = None):
+    def _change_backend(cls, engine_backend: AvailableBackends, use_pykeops: bool = False,
+                        use_gpu: bool = True, dtype: Optional[str] = None, grads:bool = False):
         cls.dtype = DEFAULT_TENSOR_DTYPE if dtype is None else dtype
         cls.dtype_obj = cls.dtype
         match engine_backend:
@@ -99,6 +102,21 @@ def _change_backend(cls, engine_backend: AvailableBackends, use_pykeops: bool =
                 cls.dtype_obj = pytorch_copy.float32 if cls.dtype == "float32" else pytorch_copy.float64
                 cls.tensor_types = pytorch_copy.Tensor
 
+                torch.set_num_threads(torch.get_num_threads())  # Use all available threads
+                cls.COMPUTE_GRADS = grads  # Store the grads setting
+                if grads is False:
+                    cls._torch_no_grad_context = torch.no_grad()
+                    cls._torch_no_grad_context.__enter__()
+                else:
+                    # If there was a previous context, exit it first
+                    if hasattr(cls, '_torch_no_grad_context') and cls._torch_no_grad_context is not None:
+                        try:
+                            cls._torch_no_grad_context.__exit__(None, None, None)
+                        except:
+                            pass  # Context might already be exited
+                    cls._torch_no_grad_context = None
+                    torch.set_grad_enabled(True)
+                    
                 cls.use_pykeops = use_pykeops  # TODO: Make this compatible with pykeops
                 if (use_pykeops):
                     import pykeops

gempy_engine/modules/evaluator/generic_evaluator.py

@@ -1,4 +1,5 @@
 import numpy as np
+import gc
 from typing import Optional
 
 from gempy_engine.core.backend_tensor import BackendTensor
@@ -57,6 +58,10 @@ def generic_evaluator(
             if gz_field is not None:
                 gz_field[slice_array] = gz_chunk  # type: ignore
 
+    # Force garbage collection every few chunks to prevent memory buildup
+    if (i + 1) % 5 == 0 or i == n_chunks - 1:
+        gc.collect()
+        
     if n_chunks > 5:
         print(f"Chunking done: {n_chunks} chunks")
 
@@ -75,7 +80,9 @@ def _eval_on(
     try:
         scalar_field = (eval_kernel.T @ weights).reshape(-1)
     except ValueError:
-        pass
+        scalar_field = None
+    
+    del eval_kernel
 
     gx_field: Optional[np.ndarray] = None
     gy_field: Optional[np.ndarray] = None
@@ -85,17 +92,21 @@ def _eval_on(
         eval_gx = yield_evaluation_grad_kernel(
             solver_input, options.kernel_options, axis=0, slice_array=slice_array
         )
+        gx_field = (eval_gx.T @ weights).reshape(-1)  # Use BEFORE deleting
+        del eval_gx  # Clean up immediately after use
+        
         eval_gy = yield_evaluation_grad_kernel(
             solver_input, options.kernel_options, axis=1, slice_array=slice_array
         )
-        gx_field = (eval_gx.T @ weights).reshape(-1)
-        gy_field = (eval_gy.T @ weights).reshape(-1)
+        gy_field = (eval_gy.T @ weights).reshape(-1)  # Use BEFORE deleting
+        del eval_gy  # Clean up immediately after use
 
         if options.number_dimensions == 3:
             eval_gz = yield_evaluation_grad_kernel(
                 solver_input, options.kernel_options, axis=2, slice_array=slice_array
             )
-            gz_field = (eval_gz.T @ weights).reshape(-1)
+            gz_field = (eval_gz.T @ weights).reshape(-1)  # Use BEFORE deleting
+            del eval_gz  # Clean up immediately after use
         elif options.number_dimensions != 2:
             raise ValueError("`number_dimensions` must be 2 or 3")