Skip to content

Commit

Permalink
fix: expose array interface for CUDA (#2327)
Browse files Browse the repository at this point in the history
* fix: expose array interface for CUDA

* test: cover other cases
  • Loading branch information
agoose77 authored and ianna committed Mar 21, 2023
1 parent 88bcb22 commit 8c828a5
Show file tree
Hide file tree
Showing 2 changed files with 38 additions and 2 deletions.
9 changes: 7 additions & 2 deletions src/awkward/index.py
Original file line number Diff line number Diff line change
Expand Up @@ -149,8 +149,13 @@ def raw(self, nplike):
def __len__(self):
return self.length

def __array__(self, dtype=None):
return self._nplike.asarray(self._data, dtype=dtype)
@property
def __cuda_array_interface__(self):
return self._data.__cuda_array_interface__

@property
def __array_interface__(self):
return self._data.__array_interface__

def __repr__(self):
return self._repr("", "", "")
Expand Down
31 changes: 31 additions & 0 deletions tests/test_2327_array_interface.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
# BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE

import numpy as np
import pytest

import awkward as ak


def test_wrap_index_cupy():
cp = pytest.importorskip("cupy")
data = cp.arange(10, dtype=cp.int64)
index = ak.index.Index64(data)
other_index = ak.index.Index64(index)
other_data = cp.asarray(other_index)
assert cp.shares_memory(data, other_data)


def test_wrap_index_numpy():
data = np.arange(10, dtype=np.int64)
index = ak.index.Index64(data)
other_index = ak.index.Index64(index)
other_data = np.asarray(other_index)
assert np.shares_memory(data, other_data)


def test_wrap_bare_list():
data = [1, 2, 3, 4, 5]
index = ak.index.Index64(data)
other_index = ak.index.Index64(index)
other_data = np.asarray(other_index)
assert other_data.tolist() == data

0 comments on commit 8c828a5

Please sign in to comment.