Repository: helmholtz-analytics/heat

Issue:
[Bug]: Random/Randn does not work in particular setting 
### What happened?

`randn` (and `rand`) produce an error in certain constellations

### Code snippet triggering the error

```shell
# run with 11 MPI-processes 

import heat as ht
A = ht.random.randn(6667,3523,dtype=ht.float64,split=1)
```


### Error message or erroneous outcome

```shell
RuntimeError: shape '[6667, 320]' is invalid for input of size 2133439
Traceback (most recent call last):
  File "developing.py", line 77, in <module>
    A = ht.random.randn(6667,3523,dtype=ht.float64,split=1)
  File "/home/hopp_fa/heat/heat/core/random.py", line 643, in randn
    normal_tensor = rand(*args, dtype=dtype, split=split, device=device, comm=comm)
  File "/home/hopp_fa/heat/heat/core/random.py", line 472, in rand
    values = __int64_to_float64(torch.stack([x_0, x_1], dim=1).flatten()[lslice]).reshape(
RuntimeError: shape '[6667, 320]' is invalid for input of size 2133439
```


### Version

1.2.x

### Python version

3.8

### PyTorch version

1.12

### MPI version

```shell
mpi4py 3.1.3
```



Discussion:
I don't know the why or the how. But I figured out where the problem is. 
It all comes down to [this one line](https://github.com/helmholtz-analytics/heat/blob/e5e8a96c28c0508df4d2865a816e40e74ac4491c/heat/core/random.py#L128)

Changing it to this fixes the issue.
```
local_elements = [total_elements.item() / shape[split] * counts[i] for i in range(size)]
```
Apparently, torch has some precision problems during division. The division operation, that takes place when the number of elements in each process is calculated, causes some problem.

**Code:**
```
import torch
print(f"Number of elements in tensor: 6667*3523={23487841}")
print(23487841/3523)
print(torch.tensor(23487841)/3523)
```

**Output:**
```
Number of elements in tensor: 6667*3523=23487841
6667.0
tensor(6666.9995)
```

This small precision error is blown up when the number of elements is huge. That's probably why the tests did not catch the bug because the number of elements were small enough.

If that's all it is with this issue, I can create a PR. But I don't know much about this part of the code to certainly say the problem is solved.
Great work! :) It looks like you found the error 👍 

As far as I understand the code the bug you found is the only error here, so if you create the PR that would be great. 

Resolving patch:
diff --git a/heat/core/random.py b/heat/core/random.py
--- a/heat/core/random.py
+++ b/heat/core/random.py
@@ -114,8 +114,8 @@ def __counter_sequence(
         raise ValueError("Shape is to big with {} elements".format(total_elements))
 
     if split is None:
-        values = torch.ceil(total_elements / 2)
-        even_end = total_elements % 2 == 0
+        values = total_elements.item() // 2 + total_elements.item() % 2
+        even_end = total_elements.item() % 2 == 0
         lslice = slice(None) if even_end else slice(None, -1)
         start = c_1
         end = start + int(values)
@@ -125,7 +125,7 @@ def __counter_sequence(
         counts, displs, _ = comm.counts_displs_shape(shape, split)
 
         # Calculate number of local elements per process
-        local_elements = [total_elements / shape[split] * counts[i] for i in range(size)]
+        local_elements = [total_elements.item() / shape[split] * counts[i] for i in range(size)]
         cum_elements = torch.cumsum(torch.tensor(local_elements, device=device.torch_device), dim=0)
 
         # Calculate the correct borders and slices


Test patch:
diff --git a/heat/core/tests/test_random.py b/heat/core/tests/test_random.py
--- a/heat/core/tests/test_random.py
+++ b/heat/core/tests/test_random.py
@@ -230,6 +230,11 @@ def test_rand(self):
         self.assertFalse(np.array_equal(a, c))
         self.assertFalse(np.array_equal(b, c))
 
+        # To check working with large number of elements
+        ht.random.randn(6667, 3523, dtype=ht.float64, split=None)
+        ht.random.randn(6667, 3523, dtype=ht.float64, split=0)
+        ht.random.randn(6667, 3523, dtype=ht.float64, split=1)
+
     def test_randint(self):
         # Checked that the random values are in the correct range
         a = ht.random.randint(low=0, high=10, size=(10, 10), dtype=ht.int64)