IntelPython / dpctl

Python SYCL bindings and SYCL-based Python Array API library
https://intelpython.github.io/dpctl/
Apache License 2.0
97 stars 29 forks source link

[Feature request] Support different `out` dtypes for elementwise func #1717

Open vlad-perevezentsev opened 1 week ago

vlad-perevezentsev commented 1 week ago

Current implementations of elementwise functions in dpctl.tensor do not support out data type other than bool and raise ValueError if out is not bool.

import dpctl.tensor as dpt

a = dpt.asarray([1,2,3])
out = dpt.empty_like(a,dtype='f4')
dpt.logical_not(a,out=out)

# ValueError: Output array of type bool is needed, got float32

While numpy can cast the result depending on out data type.

import numpy

a_np = dpt.asnumpy(a)
out_np = dpt.asnumpy(out)
numpy.logical_not(a_np, out_np)

# array([0., 0., 0.], dtype=float32)

Since the Python Array API has no description for out parameter we can implement any logic we want. I think it would be useful to follow the logic of numpy and cast the result depending on out data type.

ndgrigorian commented 1 week ago

I would like to offer a correction that not only bool is supported, it depends on the expected output type of the elementwise functions.

logical_not expects to output a boolean array, so only accepts a boolean array as out. Something like sqrt expects to output floating-point data, so only accepts floating-point data type arrays for out. The reason for this being that for the output array of the function to actually be the same array as the out array, it needs to be the appropriate data type, since we can't cast the array data in-place. In such a case, the returned array would not be the same as the user input out.