Skip to content

Commit 96636b2

Browse files
Addressed docstrings/formatting PR feedback
1 parent f276b8d commit 96636b2

File tree

1 file changed

+41
-3
lines changed

1 file changed

+41
-3
lines changed

dpctl/tensor/_dlpack.pyx

Lines changed: 41 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,10 @@ def get_build_dlpack_version():
9090
"""
9191
Returns the string value of DLPACK_VERSION from dlpack.h
9292
`dpctl.tensor` was built with.
93+
94+
Returns:
95+
A string value of the version of DLPack used to build
96+
`dpctl`.
9397
"""
9498
return str(DLPACK_VERSION)
9599

@@ -112,9 +116,26 @@ cdef void _managed_tensor_deleter(DLManagedTensor *dlm_tensor) with gil:
112116

113117
cpdef to_dlpack_capsule(usm_ndarray usm_ary) except+:
114118
"""
119+
to_dlpack_capsule(usm_ary)
120+
115121
Constructs named Python capsule object referencing
116122
instance of `DLManagerTensor` from
117123
:class:`dpctl.tensor.usm_ndarray` instance.
124+
125+
Args:
126+
usm_ary: An instance of :class:`dpctl.tensor.usm_ndarray`
127+
Returns:
128+
Python a new capsule with name "dltensor" that contains
129+
a pointer to `DLManagedTensor` struct.
130+
Raises:
131+
DLPackCreationError: when array can be represented as
132+
DLPack tensor. This may happen when array was allocated
133+
on a partitioned sycl device, or its USM allocation is
134+
not bound to the platform default SYCL context.
135+
MemoryError: when host allocation to needed for `DLManagedTensor`
136+
did not succeed.
137+
ValueError: when array elements data type could not be represented
138+
in `DLManagedTensor`.
118139
"""
119140
cdef c_dpctl.SyclQueue ary_sycl_queue
120141
cdef c_dpctl.SyclDevice ary_sycl_device
@@ -144,15 +165,13 @@ cpdef to_dlpack_capsule(usm_ndarray usm_ary) except+:
144165
"to_dlpack_capsule: DLPack can only export arrays allocated on "
145166
"non-partitioned SYCL devices."
146167
)
147-
# TODO: check that ary_sycl_context is the default context
148168
default_context = dpctl.SyclQueue(ary_sycl_device).sycl_context
149169
if not usm_ary.sycl_context == default_context:
150170
raise DLPackCreationError(
151171
"to_dlpack_capsule: DLPack can only export arrays based on USM "
152172
"allocations bound to a default platform SYCL context"
153173
)
154174

155-
156175
dlm_tensor = <DLManagedTensor *> stdlib.malloc(
157176
sizeof(DLManagedTensor))
158177
if dlm_tensor is NULL:
@@ -244,9 +263,25 @@ cdef class _DLManagedTensorOwner:
244263

245264
cpdef usm_ndarray from_dlpack_capsule(object py_caps) except +:
246265
"""
266+
from_dlpack_capsule(caps)
267+
247268
Reconstructs instance of :class:`dpctl.tensor.usm_ndarray` from
248269
named Python capsule object referencing instance of `DLManagedTensor`
249270
without copy. The instance forms a view in the memory of the tensor.
271+
272+
Args:
273+
caps: Python capsule with name "dltensor" expected to reference
274+
an instance of `DLManagedTensor` struct.
275+
Returns:
276+
Instance of :class:`dpctl.tensor.usm_ndarray` with a view into
277+
memory of the tensor. Capsule is renamed to "used_dltensor" upon
278+
success.
279+
Raises:
280+
TypeError: if argument is not a "dltensor" capsule.
281+
ValueError: if argument is "used_dltensor" capsule,
282+
if the USM pointer is not bound to the reconstructed
283+
sycl context, or the DLPack's device_type is not supported
284+
by dpctl.
250285
"""
251286
cdef DLManagedTensor *dlm_tensor = NULL
252287
cdef bytes usm_type
@@ -355,7 +390,7 @@ cpdef usm_ndarray from_dlpack_capsule(object py_caps) except +:
355390
raise ValueError(
356391
"Can not import DLPack tensor with type code {}.".format(
357392
<object>dlm_tensor.dl_tensor.dtype.code
358-
)
393+
)
359394
)
360395
res_ary = usm_ndarray(
361396
py_shape,
@@ -379,6 +414,9 @@ cpdef from_dlpack(array):
379414
object `obj` that implements `__dlpack__` protocol. The output
380415
array is always a zero-copy view of the input.
381416
417+
Args:
418+
A Python object representing an array that supports `__dlpack__`
419+
protocol.
382420
Raises:
383421
TypeError: if `obj` does not implement `__dlpack__` method.
384422
ValueError: if zero copy view can not be constructed because

0 commit comments

Comments
 (0)