Rename c_ptr to array_ptr

Resolves #170

CHANGELOG.md

@@ -27,6 +27,7 @@
 - Improved default `Figure` parameters
 - Added option to set label size in `Figure`
 - `name` no longer reserves `-n` abbreviation
+- Renamed `c_ptr` to the more sensible `array_ptr`
 
 ### Bug fixes
 
@@ -36,6 +37,7 @@
 - Fixed wrong return type in `exp_moving_avg`
 - Appending now works with `jsonl`
 - Fixed wrong type annotation in `Figure`'s `figsize` argument
+- `array_ptr` (formerly `c_ptr`) now works for CUDA tensors
 
 ## 2.0.0 - Breaking changes
 

pelutils/__init__.py

@@ -146,15 +146,12 @@ def __exit__(self, *__):
             else:
                 os.environ[var] = value
 
-def c_ptr(arr: np.ndarray | torch.Tensor) -> ctypes.c_void_p:
-    """
-    Returns a C pointer that can be used to import a contiguous
-    numpy array or torch tensor into a C/C++/Rust function. This
-    function is mostly useful when not using Python's C api and
-    instead interfacing with .so files directly.
-    """
+def array_ptr(arr: np.ndarray | torch.Tensor) -> ctypes.c_void_p:
+    """ Returns a pointer to a numpy array or torch tensor which can be used to interact
+    with it in low-level languages like C/C++/Rust. This function is mostly useful
+    when not using Python's C api and instead interfacing with .so files directly. """
     if _has_torch and isinstance(arr, torch.Tensor):
-        arr = arr.numpy()
+        return ctypes.c_void_p(arr.data_ptr())
     if not isinstance(arr, np.ndarray):
         raise TypeError(f"Array should be of type np.ndarray or torch.Tensor, not {type(arr)}")
     if not arr.flags.c_contiguous:
@@ -166,9 +163,8 @@ def split_path(path: str) -> list[str]:
     return os.path.normpath(path).split(os.sep)
 
 def binary_search(element: _T, iterable: Sequence[_T], *, _start=0, _end=-1) -> int | None:
-    """ Get the index of element in sequence using binary search
-    Assumes iterable is sorted in ascending order
-    Returns None if the element is not found """
+    """ Get the index of element in sequence using binary search. Assumes iterable is
+    sorted in ascending order. Returns None if the element is not found. """
     if _end == -1:  # Entered on first call
         _end = len(iterable)
         # Make sure element actually exists in array

tests/test_init.py

@@ -10,7 +10,7 @@
 import torch
 
 from pelutils import EnvVars, UnsupportedOS, reverse_line_iterator, except_keys,\
-    split_path, binary_search, raises, thousands_seperators, OS, c_ptr,\
+    split_path, binary_search, raises, thousands_seperators, OS, array_ptr,\
     get_timestamp, get_timestamp_for_files, HardwareInfo
 from pelutils.tests import UnitTestCollection
 
@@ -146,12 +146,14 @@ def test_except_keys(self):
         assert "a" in d and "b" in d
         assert "a" in d2 and "b" not in d2
 
-    def test_c_ptr(self):
+    def test_array_ptr(self):
         with pytest.raises(TypeError):
-            c_ptr(None)
+            array_ptr(None)
         with pytest.raises(ValueError):
-            c_ptr(np.arange(5)[::2])
-        assert isinstance(c_ptr(torch.arange(5)), ctypes.c_void_p)
+            array_ptr(np.arange(5)[::2])
+        assert isinstance(array_ptr(torch.arange(5)), ctypes.c_void_p)
+        a = torch.arange(5)
+        assert array_ptr(a).value == array_ptr(a.numpy()).value
 
     def test_get_timestamp(self):
         for date in False, True: