@@ -271,25 +271,28 @@ cpdef utils.dpnp_descriptor dpnp_ones_like(result_shape, result_dtype):
271
271
return call_fptr_1out(DPNP_FN_ONES_LIKE, utils._object_to_tuple(result_shape), result_dtype)
272
272
273
273
274
- cpdef dparray dpnp_trace(arr, offset = 0 , axis1 = 0 , axis2 = 1 , dtype = None , out = None ):
274
+ cpdef utils.dpnp_descriptor dpnp_trace(utils.dpnp_descriptor arr, offset = 0 , axis1 = 0 , axis2 = 1 , dtype = None , out = None ):
275
275
if dtype is None :
276
276
dtype_ = arr.dtype
277
277
else :
278
278
dtype_ = dtype
279
279
280
280
cdef dparray diagonal_arr = dpnp.diagonal(arr, offset, axis1, axis2)
281
+ cdef size_t diagonal_ndim = diagonal_arr.ndim
282
+ cdef dparray_shape_type diagonal_shape = diagonal_arr.shape
281
283
282
284
cdef DPNPFuncType param1_type = dpnp_dtype_to_DPNPFuncType(arr.dtype)
283
285
cdef DPNPFuncType param2_type = dpnp_dtype_to_DPNPFuncType(dtype_)
284
286
285
287
cdef DPNPFuncData kernel_data = get_dpnp_function_ptr(DPNP_FN_TRACE, param1_type, param2_type)
286
288
287
- result_type = dpnp_DPNPFuncType_to_dtype(< size_t > kernel_data.return_type)
288
- cdef dparray result = dparray(diagonal_arr.shape[:- 1 ], dtype = result_type)
289
+ # ceate result array with type given by FPTR data
290
+ cdef dparray_shape_type result_shape = diagonal_shape[:- 1 ]
291
+ cdef utils.dpnp_descriptor result = utils.create_output_descriptor(result_shape, kernel_data.return_type, None )
289
292
290
293
cdef fptr_dpnp_trace_t func = < fptr_dpnp_trace_t > kernel_data.ptr
291
294
292
- func(diagonal_arr.get_data(), result.get_data(), < size_t * > diagonal_arr._dparray_shape. data(), diagonal_arr.ndim )
295
+ func(diagonal_arr.get_data(), result.get_data(), < size_t * > diagonal_shape. data(), diagonal_ndim )
293
296
294
297
return result
295
298
0 commit comments