helmholtz-analytics / heat

Distributed tensors and Machine Learning framework with GPU and MPI acceleration in Python
https://heat.readthedocs.io/
MIT License
209 stars 53 forks source link

[Bug]: `_like` factory functions mistakenly set device to Default #1464

Closed ClaudiaComito closed 4 months ago

ClaudiaComito commented 5 months ago

What happened?

Functions ht.empty_like(a), ht.full_like(a), ht.ones_like(a) and ht.zeros_like(a) return new DNDarrays with the same type, shape and data distribution of a but not with the same device. The device is mistakenly set to the globally-set default, this explains why the CI didn't catch it, as @mrfh92 correctly pointed out.

The equivalent torch functions return new tensors with the same attributes including the same device as the input tensor.

Code snippet triggering the error

import heat as ht

a = ht.arange(10, device="gpu")
b = ht.ones_like(a)
a + b

Error message or erroneous outcome

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

Version

main (development branch)

Python version

3.11

PyTorch version

2.1

MPI version

No response