XuPy 1.7.0__tar.gz → 1.7.1__tar.gz

{xupy-1.7.0 → xupy-1.7.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: XuPy
-Version: 1.7.0
+Version: 1.7.1
 Summary: GPU Accelerated masked arrays with automatic handling of CPU and GPU arrays.
 Author: Pietro Ferraiuolo
 Author-email: Pietro Ferraiuolo <pietro.ferraiuolo@inaf.it>

{xupy-1.7.0 → xupy-1.7.1}/XuPy.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: XuPy
-Version: 1.7.0
+Version: 1.7.1
 Summary: GPU Accelerated masked arrays with automatic handling of CPU and GPU arrays.
 Author: Pietro Ferraiuolo
 Author-email: Pietro Ferraiuolo <pietro.ferraiuolo@inaf.it>

{xupy-1.7.0 → xupy-1.7.1}/XuPy.egg-info/SOURCES.txt

@@ -9,6 +9,7 @@ XuPy.egg-info/requires.txt
 XuPy.egg-info/top_level.txt
 test/test_arithmetic_compatibility.py
 test/test_core.py
+test/test_cpu_memory_context.py
 test/test_large_array_printing.py
 test/test_ma_core.py
 test/test_ma_extras.py

{xupy-1.7.0 → xupy-1.7.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "XuPy"
-version = "1.7.0"
+version = "1.7.1"
 description = "GPU Accelerated masked arrays with automatic handling of CPU and GPU arrays."
 authors = [
     { name = "Pietro Ferraiuolo", email = "pietro.ferraiuolo@inaf.it" }

{xupy-1.7.0 → xupy-1.7.1}/test/test_core.py

@@ -19,7 +19,7 @@ except ImportError:
     HAS_CUPY = False
 
 import xupy as xp
-from xupy._core import NumpyContext, on_gpu
+from xupy._core import NumpyContext, _CPUMemoryContext, on_gpu
 
 # Skip all tests if CuPy is not available
 pytestmark = pytest.mark.skipif(not HAS_CUPY, reason="CuPy not available")

xupy-1.7.1/test/test_cpu_memory_context.py

@@ -0,0 +1,123 @@
+"""
+Tests for _CPUMemoryContext.
+
+These tests do NOT require CuPy/CUDA and run purely on CPU (NumPy).
+They validate that the CPU memory context manager exposes the same
+interface as the GPU MemoryContext and behaves as a well-behaved
+no-op context on CPU.
+"""
+import pytest
+import numpy as np
+
+import xupy as xp
+from xupy._core import _CPUMemoryContext, on_gpu
+
+
+class TestCPUMemoryContext:
+    """Tests for _CPUMemoryContext (always available, no GPU required)."""
+
+    def test_cpu_memory_context_is_importable(self):
+        """_CPUMemoryContext can be imported from _core and is a class."""
+        assert callable(_CPUMemoryContext)
+
+    def test_cpu_memory_context_enter_returns_self(self):
+        """__enter__ returns the context object itself."""
+        ctx = _CPUMemoryContext()
+        with ctx as c:
+            assert c is ctx
+
+    def test_cpu_memory_context_basic_usage(self, capsys):
+        """Context manager completes without raising and prints a summary."""
+        with _CPUMemoryContext() as ctx:
+            arr = np.array([1, 2, 3])
+            assert arr is not None
+        captured = capsys.readouterr()
+        assert "Session completed" in captured.out
+
+    def test_cpu_memory_context_default_parameters(self):
+        """Default parameters match the GPU MemoryContext API."""
+        ctx = _CPUMemoryContext()
+        assert ctx.device_id is None
+        assert ctx.auto_cleanup is True
+        assert ctx.memory_threshold == 0.9
+        assert ctx.monitor_interval == 1.0
+
+    def test_cpu_memory_context_custom_parameters(self):
+        """Custom parameters are stored correctly."""
+        ctx = _CPUMemoryContext(
+            device_id=0,
+            auto_cleanup=False,
+            memory_threshold=0.8,
+            monitor_interval=2.0,
+        )
+        assert ctx.device_id == 0
+        assert ctx.auto_cleanup is False
+        assert ctx.memory_threshold == 0.8
+        assert ctx.monitor_interval == 2.0
+
+    def test_cpu_memory_context_get_memory_info_returns_dict(self):
+        """get_memory_info always returns a dict with at least 'device'."""
+        ctx = _CPUMemoryContext()
+        info = ctx.get_memory_info()
+        assert isinstance(info, dict)
+        assert "device" in info
+        assert info["device"] == "cpu"
+
+    def test_cpu_memory_context_get_memory_info_keys(self):
+        """get_memory_info returns expected keys when psutil is available."""
+        try:
+            import psutil  # noqa: F401
+            psutil_available = True
+        except ImportError:
+            psutil_available = False
+
+        ctx = _CPUMemoryContext()
+        info = ctx.get_memory_info()
+        if psutil_available:
+            for key in ("total", "free", "used", "memory_percent"):
+                assert key in info, f"Missing key '{key}' in memory info"
+            assert info["total"] > 0
+            assert 0.0 <= info["memory_percent"] <= 1.0
+        else:
+            assert "error" in info
+
+    def test_cpu_memory_context_no_op_methods(self):
+        """All GPU-specific methods are no-ops (do not raise)."""
+        ctx = _CPUMemoryContext()
+        ctx.track_object(object())
+        ctx.clear_cache()
+        ctx.aggressive_cleanup()
+        ctx.emergency_cleanup()
+        ctx.auto_cleanup_if_needed()
+        ctx.monitor_memory(duration=0.0)
+        ctx.force_memory_deallocation()
+        ctx.force_memory_pool_reset()
+
+    def test_cpu_memory_context_check_memory_pressure(self):
+        """check_memory_pressure returns a bool."""
+        ctx = _CPUMemoryContext()
+        result = ctx.check_memory_pressure()
+        assert isinstance(result, bool)
+
+    def test_cpu_memory_context_repr(self):
+        """__repr__ contains 'MemoryContext' and 'cpu'."""
+        ctx = _CPUMemoryContext()
+        r = repr(ctx)
+        assert "MemoryContext" in r
+        assert "cpu" in r
+
+    def test_cpu_memory_context_exception_propagation(self):
+        """Exceptions raised inside the context propagate normally."""
+        with pytest.raises(ValueError):
+            with _CPUMemoryContext():
+                raise ValueError("test error")
+
+    def test_cpu_memory_context_exposed_when_on_cpu(self):
+        """When running in CPU mode, xp.MemoryContext is _CPUMemoryContext."""
+        if not on_gpu:
+            assert hasattr(xp, "MemoryContext")
+            assert xp.MemoryContext is _CPUMemoryContext
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

xupy-1.7.1/xupy/__version__.py

@@ -0,0 +1 @@
+__version__ = "1.7.1"

{xupy-1.7.0 → xupy-1.7.1}/xupy/_core.py

@@ -11,10 +11,12 @@ import time as _time
 import builtins as _b
 import sys as _sys
 from . import typings as _t
+from contextlib import contextmanager as _contextmanager
 from ._cupy_install import __check_availability__ as __check__
 
 _GPU = False
 _GPU_AVAILABLE = False
+_MULTIGPU = False
 
 __check__.xupy_init()
 __cuda_version__ = __check__.get_cuda_version()
@@ -25,8 +27,10 @@ try:
     import cupy as _xp # type: ignore
 
     _B2mb_ = 1024 * 1000  # using MB = 1,000,000 bytes
+    _Btgb_ = 1024 * 1000 * 1000  # using GB = 1,000,000,000 bytes
     n_gpus = _xp.cuda.runtime.getDeviceCount()
     if n_gpus > 1:
+        _MULTIGPU = True
         gpus = {}
         line1 = """
 [XuPy] Multiple GPUs detected:
@@ -68,6 +72,7 @@ except Exception as err:
     from numpy import *  # type: ignore
 
 on_gpu = _GPU
+has_multi_gpu = _MULTIGPU
 
 # Capture the set of names brought in by the wildcard import
 _backend_names = frozenset(
@@ -77,6 +82,48 @@ _backend_names = frozenset(
 
 _mode_names: set[str] = set()
 
+def _array_size(
+    shape: tuple[int] | list[tuple[int]],
+    dtype: _t.DTypeLike = _xp.float32,
+    out_unit: str = 'MB',
+) -> int:
+    """
+    Computes the expected allocated size on GPU of an array with size `shape` 
+    and data type `dtype`.
+
+    Parameters
+    ----------
+    shape : tuple[int] | list[tuple[int]]
+        The shape of the array. Can input multiple shapes as a list, and the
+        result will be the total size of all arrays combined.
+    dtype : DTypeLike, optional
+        The data type of the array elements (default: float32).
+    out_unit : str, optional
+        The unit for the output size. Can be 'MB' or 'GB' (default: 'MB').
+
+    Returns
+    -------
+    size : int
+        The size of the array in the specified unit.
+
+    Examples
+    --------
+    >>> import xupy as xp
+    >>> arr = xp.array([1, 2, 3])
+    >>> xp.array_size(arr)
+    12  # 3 elements * 4 bytes per int32
+    """
+    norm = _B2mb_ if out_unit == 'MB' else _Btgb_
+    if isinstance(shape, tuple):
+        if isinstance(shape[0], int):
+            shape = [shape]  # single shape case
+    size = []
+    for s in shape:
+        itemsize = _np.dtype(dtype).itemsize
+        num_elements = _np.prod(s)
+        size_bytes = num_elements * itemsize
+        size.append(int(size_bytes / norm))
+    return int(_np.sum(size))
 
 # --- NUMPY Context manager ---
 class NumpyContext:
@@ -114,16 +161,282 @@ class NumpyContext:
     def __repr__(self) -> str:
         """String representation of the context manager."""
         if _GPU:
+
             return f"NumpyContext(original_device={self.original_device})"
         else:
             return "NumpyContext(no_gpu=True)"
 
 
+# ---------------------------------------------------------------------------
+# CPU Memory Context Manager (always available, no-op mock for CPU mode)
+# ---------------------------------------------------------------------------
+
+class _CPUMemoryContext:
+    """CPU memory context manager — a no-op counterpart to the GPU _MemoryContext.
+
+    Provides the same interface as the GPU ``MemoryContext`` so that code written
+    against ``xp.MemoryContext`` runs transparently on CPU (NumPy) without any
+    changes.  All GPU-specific operations (pool cleanup, device synchronisation,
+    etc.) are silently skipped.
+
+    Example
+    -------
+    >>> import xupy as xp          # running in CPU mode
+    >>> with xp.MemoryContext() as ctx:
+    ...     arr = xp.array([1, 2, 3])
+    ...     print(ctx.get_memory_info())
+    """
+
+    def __init__(
+        self,
+        device_id: _t.Optional[int] = None,
+        auto_cleanup: bool = True,
+        force_cleanup: bool = False,
+        print_report: bool = True,
+        memory_threshold: float = 0.9,
+        monitor_interval: float = 1.0,
+    ):
+        """
+        Parameters
+        ----------
+        device_id : int, optional
+            Ignored on CPU; present for API compatibility with the GPU version.
+        auto_cleanup : bool, optional
+            Kept for API compatibility; no cleanup is performed on CPU.
+        memory_threshold : float, optional
+            Kept for API compatibility; no threshold enforcement on CPU.
+        monitor_interval : float, optional
+            Kept for API compatibility; no monitoring is performed on CPU.
+        """
+        self.device_id = device_id
+        self.auto_cleanup = auto_cleanup
+        self.force_cleanup = force_cleanup
+        self.print_report = print_report
+        self.memory_threshold = memory_threshold
+        self.monitor_interval = monitor_interval
+
+        self._start_time: _t.Optional[float] = None
+
+    def __enter__(self):
+        """Enter the CPU memory context."""
+        self._start_time = _time.time()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Exit the CPU memory context (no-op cleanup)."""
+        if self._start_time is not None:
+            duration = _time.time() - self._start_time
+            if self.print_report:
+                print(f"[MemoryContext] Session completed in {duration:.2f}s (CPU mode)")
+
+    def track_object(self, obj):
+        """No-op: object tracking is not needed on CPU."""
+        pass
+
+    def clear_cache(self):
+        """No-op: no GPU memory pool to clear on CPU."""
+        pass
+
+    def aggressive_cleanup(self):
+        """No-op: no GPU memory to aggressively free on CPU."""
+        pass
+
+    def emergency_cleanup(self):
+        """No-op: no GPU memory emergency cleanup needed on CPU."""
+        pass
+
+    def get_memory_info(self) -> dict:
+        """Return basic CPU/RAM memory information where available.
+
+        Uses ``psutil`` when installed; otherwise returns a minimal dict.
+        """
+        info: dict = {"device": "cpu"}
+        try:
+            import psutil  # type: ignore
+            vm = psutil.virtual_memory()
+            info.update(
+                {
+                    "total": vm.total,
+                    "free": vm.available,
+                    "used": vm.used,
+                    "memory_percent": vm.percent / 100.0,
+                }
+            )
+        except ImportError:
+            info["error"] = "psutil not installed; install it for detailed CPU memory info"
+        return info
+
+    def check_memory_pressure(self) -> bool:
+        """Check if RAM usage is above the threshold (requires psutil)."""
+        mem_info = self.get_memory_info()
+        if "memory_percent" in mem_info:
+            return mem_info["memory_percent"] > self.memory_threshold
+        return False
+
+    def auto_cleanup_if_needed(self):
+        """No-op: no GPU pressure-based cleanup on CPU."""
+        pass
+
+    def monitor_memory(self, duration: float = 10.0):
+        """No-op: memory monitoring is not performed in CPU mode."""
+        pass
+
+    def force_memory_deallocation(self):
+        """No-op: forced GPU memory deallocation is not applicable on CPU."""
+        pass
+
+    def force_memory_pool_reset(self):
+        """No-op: GPU memory pool reset is not applicable on CPU."""
+        pass
+
+    def __repr__(self) -> str:
+        """String representation of the CPU memory context."""
+        mem_info = self.get_memory_info()
+        if "used" in mem_info:
+            used_mb = mem_info["used"] / (1024 * 1000)
+            total_mb = mem_info["total"] / (1024 * 1000)
+            percent = mem_info["memory_percent"] * 100
+            return f"MemoryContext(device=cpu, memory={used_mb:.2f}/{total_mb:.2f} MB ({percent:.1f}%))"
+        return "MemoryContext(device=cpu)"
+
+
 # ---------------------------------------------------------------------------
 # GPU-only definitions (only created when CuPy was successfully loaded)
 # ---------------------------------------------------------------------------
 
 if _GPU_AVAILABLE:
+    
+    def _allocate_with_fallback_device(
+        array: _t.ArrayLike,
+        dtype: _t.Optional[_t.DTypeLike] = None,
+        preferred_device: _t.Optional[int] = None,
+        safety_factor: float = 1.10,
+        reserve_mb: int = 128,
+        ) -> tuple[_t.NDArray[_t.Any], int]:
+        """
+        Allocate an array on the current/preferred GPU if enough memory is available.
+        If not, try other GPUs when multi-GPU is available.
+        
+        Parameters
+        ----------
+        array : ArrayLike
+            Input data to allocate on GPU.
+        dtype : DTypeLike, optional
+            Target dtype for allocation. If None, uses input dtype when available.
+        preferred_device : int, optional
+            Device to try first. If None, uses current CUDA device.
+        safety_factor : float, optional
+            Extra multiplicative margin on top of estimated size.
+        reserve_mb : int, optional
+            Additional fixed memory cushion to reduce OOM risk.
+
+        Returns
+        -------
+        gpu_array : NDArray
+            Allocated array on the selected GPU.
+        device_id : int
+            GPU id where the allocation was performed.
+
+        Raises
+        ------
+        RuntimeError
+            If GPU backend is not available.
+        MemoryError
+            If no device has enough free memory.
+        """
+        if not _GPU_AVAILABLE:
+            raise RuntimeError("[XuPy] GPU backend is not available.")
+
+        # Keep behavior explicit: this helper is for GPU allocation.
+        if not on_gpu:
+            raise RuntimeError("[XuPy] XuPy is in CPU mode. Call use_gpu() first.")
+
+        # Resolve dtype used for memory estimate and allocation.
+        target_dtype = dtype if dtype is not None else getattr(array, "dtype", _xp.float32)
+
+        # Estimate required memory in MB using existing XuPy helper.
+        shape = getattr(array, "shape", None)
+        if shape is None:
+            arr_np = _np.asarray(array, dtype=target_dtype)
+            shape = arr_np.shape
+        required_mb = _array_size(tuple(shape), dtype=target_dtype, out_unit="MB")
+        required_mb = int(required_mb * safety_factor) + int(reserve_mb)
+
+        current_device = _get_device()
+        first_device = current_device if preferred_device is None else int(preferred_device)
+
+        # Device probe order: preferred/current first, then the others.
+        device_order = [first_device]
+        if has_multi_gpu:
+            device_order.extend([d for d in range(n_gpus) if d != first_device])
+
+        for dev_id in device_order:
+            try:
+                with _on_device(dev_id):
+                    free_b, _ = _xp.cuda.runtime.memGetInfo()
+                    free_mb = int(free_b / _B2mb_)
+
+                if free_mb >= required_mb:
+                    with _on_device(dev_id):
+                        gpu_array = _xp.asarray(array, dtype=target_dtype)
+                    return gpu_array, dev_id
+            except Exception:
+                # Skip unavailable/busy devices and continue probing.
+                continue
+
+        raise MemoryError(
+            f"[XuPy] Cannot allocate array (~{required_mb} MB incl. margin) "
+            f"on probed devices {device_order}."
+        )
+
+
+    @_contextmanager
+    def _on_device(device_id: int):
+        """
+        Context manager to temporarily set the CUDA device for computations (cupy).
+
+        Parameters
+        ----------
+        device_id : int
+            The ID of the CUDA device to set as default within the context.
+            
+            If ``-1`` is passed, it will switch to CPU mode within the context 
+            and restore GPU mode on exit.
+
+        Raises
+        ------
+        RuntimeError : If the device cannot be set or if the device is already 
+            the current device.
+
+        Examples
+        --------
+        .. code-block:: python
+            with xp.on_device(0):
+                # computations here will use device 0
+                array = xp.array([1,2,3]) # array allocated on GPU 0
+            with xp.on_device(1):
+                # computations here will use device 1
+                array = xp.array([4,5,6]) # array allocated on GPU 1
+            with xp.on_device(-1):
+                # computations here will use CPU
+                array = xp.array([7,8,9]) # array allocated on CPU
+        """
+        original_device = _xp.cuda.runtime.getDevice()
+        try:
+            if device_id == -1:
+                use_cpu()
+            else:
+                _set_device(device_id)
+            yield
+        finally:
+            # Restore original device
+            try:
+                if device_id == -1:
+                    use_gpu()
+                _xp.cuda.runtime.setDevice(original_device)
+            except Exception as e:
+                print(f"Warning: Could not restore original device: {e}")
+    
 
     def _set_device(device_id: int) -> None:
         """
@@ -157,6 +470,28 @@ if _GPU_AVAILABLE:
             warnings.warn(
                 f"[XuPy] Device {device_id} is already the current device", UserWarning
             )
+    
+    def _get_device() -> int:
+        """
+        Get the current CUDA device ID.
+
+        Returns
+        -------
+        int
+            The ID of the current CUDA device.
+
+        Raises
+        ------
+        RuntimeError : If the GPU backend is not available.
+
+        Examples
+        --------
+        >>> current_device = xp.get_device()
+        >>> print(f"Current device ID: {current_device}")
+        """
+        if not _GPU_AVAILABLE:
+            return -1  # Indicate no GPU available
+        return int(_xp.cuda.runtime.getDevice())
 
     # --- GPU Memory Management Context Manager ---
     class _MemoryContext:
@@ -176,6 +511,8 @@ if _GPU_AVAILABLE:
             self,
             device_id: _t.Optional[int] = None,
             auto_cleanup: bool = True,
+            force_cleanup: bool = False,
+            print_report: bool = True,
             memory_threshold: float = 0.9,
             monitor_interval: float = 1.0,
         ):
@@ -188,15 +525,21 @@ if _GPU_AVAILABLE:
                 GPU device ID to manage. If None, uses current device.
             auto_cleanup : bool, optional
                 Whether to automatically cleanup memory on exit (default: True).
+            force_cleanup : bool, optional
+                Whether to force memory cleanup on exit (default: False).
+            print_report : bool, optional
+                Whether to print memory usage report on exit (default: True).
             memory_threshold : float, optional
                 Memory usage threshold (0-1) for automatic cleanup (default: 0.9).
             monitor_interval : float, optional
-                Interval in seconds for memory monitoring (default: 1.0).
+                Interval in milliseconds for memory monitoring (default: 1.0).
             """
             self.device_id = device_id
             self.auto_cleanup = auto_cleanup
+            self.force_cleanup = force_cleanup
             self.memory_threshold = memory_threshold
-            self.monitor_interval = monitor_interval
+            self.monitor_interval = monitor_interval/1000
+            self._print_report = print_report
 
             self._device_ctx = None
             self._original_device = None
@@ -237,33 +580,34 @@ if _GPU_AVAILABLE:
         def __exit__(self, exc_type, exc_val, exc_tb):
             """Exit the memory context with cleanup."""
             try:
-                if self.auto_cleanup:
+                if self.auto_cleanup or self.force_cleanup:
                     self.aggressive_cleanup()
 
-                # Cleanup tracked GPU objects
-                self._cleanup_gpu_objects()
+                    # Cleanup tracked GPU objects
+                    self._cleanup_gpu_objects()
 
-                # Restore original device
-                if _GPU and self._device_ctx is not None:
-                    try:
-                        self._device_ctx.__exit__(exc_type, exc_val, exc_tb)
-                    except Exception as e:
-                        print(f"Warning: Error restoring device context: {e}")
+                    # Restore original device
+                    if _GPU and self._device_ctx is not None:
+                        try:
+                            self._device_ctx.__exit__(exc_type, exc_val, exc_tb)
+                        except Exception as e:
+                            print(f"Warning: Error restoring device context: {e}")
 
                 # Final memory report
-                if self._start_time:
-                    duration = _time.time() - self._start_time
-                    final_mem = self.get_memory_info()
-                    if "used" in final_mem:
-                        memory_delta = final_mem["used"] - self._initial_memory
-                        print(f"[MemoryContext] Session completed in {duration:.2f}s")
-                        print(
-                            f"[MemoryContext] Memory delta: {memory_delta / (_B2mb_):.2f} MB"
-                        )
-                        if self._cleanup_count > 0:
+                if self._print_report:
+                    if self._start_time:
+                        duration = _time.time() - self._start_time
+                        final_mem = self.get_memory_info()
+                        if "used" in final_mem:
+                            memory_delta = final_mem["used"] - self._initial_memory
+                            print(f"[MemoryContext] Session completed in {duration:.2f}s")
                             print(
-                                f"[MemoryContext] Cleanup operations: {self._cleanup_count}"
+                                f"[MemoryContext] Memory delta: {memory_delta / (_B2mb_):.2f} MB"
                             )
+                            if self._cleanup_count > 0:
+                                print(
+                                    f"[MemoryContext] Cleanup operations: {self._cleanup_count}"
+                                )
 
             except Exception as e:
                 print(f"Warning: Error during memory context cleanup: {e}")
@@ -322,7 +666,8 @@ if _GPU_AVAILABLE:
             if not _GPU:
                 return
 
-            print("[MemoryContext] Performing aggressive memory cleanup...")
+            if self._print_report:
+                print("[MemoryContext] Performing aggressive memory cleanup...")
             self._cleanup_count += 1
 
             # Force garbage collection
@@ -385,17 +730,18 @@ if _GPU_AVAILABLE:
             except Exception as e:
                 print(f"Warning: Could not get CUDA memory info: {e}")
 
-            # As a last resort, try memory pool reset
-            try:
-                self.force_memory_pool_reset()
-            except Exception as e:
-                print(f"Warning: Memory pool reset failed: {e}")
+            if self.force_cleanup:
+                # As a last resort, try memory pool reset
+                try:
+                    self.force_memory_pool_reset()
+                except Exception as e:
+                    print(f"Warning: Memory pool reset failed: {e}")
 
-            # Final attempt: force memory deallocation
-            try:
-                self.force_memory_deallocation()
-            except Exception as e:
-                print(f"Warning: Forced memory deallocation failed: {e}")
+                # Final attempt: force memory deallocation
+                try:
+                    self.force_memory_deallocation()
+                except Exception as e:
+                    print(f"Warning: Forced memory deallocation failed: {e}")
 
         def emergency_cleanup(self):
             """Emergency cleanup for out-of-memory situations."""
@@ -495,9 +841,9 @@ if _GPU_AVAILABLE:
 
                 info = {
                     "device": int(device_to_query),
-                    "total": int(total),
-                    "free": int(free),
-                    "used": int(used),
+                    "total": int(total / _B2mb_),
+                    "free": int(free / _B2mb_),
+                    "used": int(used / _B2mb_),
                     "memory_percent": memory_percent,
                     "pool_used": pool_used,
                     "pool_capacity": pool_capacity,
@@ -612,10 +958,8 @@ if _GPU_AVAILABLE:
                 # Check memory after
                 free_after, _ = _xp.cuda.runtime.memGetInfo()
                 freed = free_after - free_before
-                # print(
-                #     f"[MemoryContext] Memory after forced deallocation: {free_after/(_B2mb_):.2f}/{total/(_B2mb_):.2f} MB"
-                # )
-                print(f"[MemoryContext] Memory freed: {freed/(_B2mb_):.2f} MB")
+                if self._print_report:
+                    print(f"[MemoryContext] Memory freed: {freed/(_B2mb_):.2f} MB")
 
             except Exception as e:
                 print(f"Warning: Could not force memory deallocation: {e}")
@@ -625,7 +969,8 @@ if _GPU_AVAILABLE:
             if not _GPU:
                 return
 
-            print("[MemoryContext] Performing memory pool reset...")
+            if self._print_report:
+                print("[MemoryContext] Performing memory pool reset...")
             try:
                 # Get current pool
                 old_pool = _xp.get_default_memory_pool()
@@ -647,7 +992,8 @@ if _GPU_AVAILABLE:
                 # Synchronize to ensure operations are complete
                 _xp.cuda.runtime.deviceSynchronize()
 
-                print("[MemoryContext] Memory pool reset completed")
+                if self._print_report:
+                    print("[MemoryContext] Memory pool reset completed")
 
             except Exception as e:
                 print(f"Warning: Could not reset memory pool: {e}")
@@ -675,6 +1021,7 @@ if _GPU_AVAILABLE:
 
 def _gpu_definitions() -> dict:
     """Return dict of names exposed only in GPU mode."""
+
     def asmarray(array: _t.NDArray[_t.Any]) -> _t.MaskedArray:
         """
         Converts an object to a masked array on the GPU.
@@ -699,6 +1046,8 @@ def _gpu_definitions() -> dict:
         'npma': _np.ma,
         'asmarray': asmarray,
         'set_device': _set_device,
+        'array_size': _array_size,
+        'on_device': _on_device,
         'MemoryContext': _MemoryContext,
     }
 
@@ -726,8 +1075,11 @@ def _cpu_definitions() -> dict:
         'double': _np.float64,
         'cfloat': _np.complex128,
         'cdouble': _np.complex128,
+        'array_size': _array_size,
         'asnumpy': asnumpy,
         'asmarray': asmarray,
+        'MemoryContext': _CPUMemoryContext,
+        'on_device': lambda device_id: None,  # No-op on CPU
     }
 
 

xupy-1.7.0/xupy/__version__.py

@@ -1 +0,0 @@
-__version__ = "1.7.0"