AntoinePrv commented on code in PR #51122:
URL: https://github.com/apache/arrow/pull/51122#discussion_r3959453291
##########
python/pyarrow/tests/test_dlpack.py:
##########
@@ -374,3 +374,101 @@ def test_dlpack_cuda_not_supported():
with pytest.raises(NotImplementedError, match="DLPack support is
implemented "
"only for buffers on CPU device."):
carr.__dlpack_device__()
+
+
+@requires_numpy_version("2.1.0")
+@check_bytes_allocated
[email protected]('np_type',
+ [np.uint8, np.uint16, np.uint32, np.uint64,
+ np.int8, np.int16, np.int32, np.int64,
+ np.float16, np.float32, np.float64])
+def test_tensor_from_dlpack(np_type):
+ def make_array():
+ base = np.arange(24, dtype=np_type).reshape((4, 6))
+ array = base[::2, 1::2]
+ assert not array.flags['C_CONTIGUOUS']
+ return array
+
+ # Non-contiguous, strided slice: DLPack carries explicit strides, so this
+ # should not need a copy on export.
+ tensor = pa.Tensor.from_dlpack(make_array())
+ assert isinstance(tensor, pa.Tensor)
+ gc.collect() # Attempts to free input array memory
+ np.testing.assert_array_equal(tensor.to_numpy(), make_array(), strict=True)
+
+
+@requires_numpy_version("2.1.0")
+@check_bytes_allocated
[email protected]('np_type',
+ [np.uint8, np.uint16, np.uint32, np.uint64,
+ np.int8, np.int16, np.int32, np.int64,
+ np.float16, np.float32, np.float64])
+def test_array_from_dlpack(np_type):
+ expected = np.array([1, 2, 3, 4, 5], dtype=np_type)
+ arr = pa.Array.from_dlpack(expected)
+ arr.validate(full=True)
+ assert isinstance(arr, pa.Array)
+ np.testing.assert_array_equal(arr.to_numpy(), expected, strict=True)
+
+
+@requires_numpy_version("2.1.0")
+@check_bytes_allocated
[email protected]('np_type',
+ [np.uint8, np.uint16, np.uint32, np.uint64,
+ np.int8, np.int16, np.int32, np.int64,
+ np.float16, np.float32, np.float64])
+def test_fixed_shape_tensor_array_from_dlpack(np_type):
+ source = np.arange(12, dtype=np_type).reshape((3, 2, 2))
+ arr = pa.FixedShapeTensorArray.from_dlpack(source)
+ arr.validate(full=True)
+ assert arr.type == pa.fixed_shape_tensor(pa.from_numpy_dtype(np_type), [2,
2])
+ assert arr.to_pylist() == [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]]
Review Comment:
Note that the underlying `FixedShapeTensorArray::FromTensor` predates this
PR.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
