Repository: iree-org/iree-turbine (language: python)

Issue:
Adding 0d tensors to parameter archive is unreliable
In general, getting a raw buffer representing a tensor is done like this (to add to a parameter archive):

```
flat_array = tensor.detach().flatten().contiguous().cpu().view(torch.uint8)
host_array = flat_array.numpy()
```

There is some kind of internal path in torch that is different, though, for scalar 0d tensors. It has been observed that taking a view like this seems to not always point to a persistent version of that backing data. Such tensors act a little bit differently indicating some special case optimization or something for them (i.e. if viewing them as uint8, you can directly access the contents by casting to bytes without first going to numpy()).

Through some trial and error, I have found that removing the view and directly capturing a copy as bytes seems to be reliable. While not satisfied without having run down the root cause, this seems robust.

Example:

```
        if len(tensor.shape) == 0:
            flat_array = tensor.detach().cpu().numpy()
            host_array = bytes(flat_array)
        else:
            flat_array = tensor.detach().flatten().contiguous().cpu().view(torch.uint8)
            host_array = flat_array.numpy()
```



Discussion / hints:

Fix patch:
diff --git a/shark_turbine/aot/params.py b/shark_turbine/aot/params.py
index 63bcc1b0..6ba0ecbe 100644
--- a/shark_turbine/aot/params.py
+++ b/shark_turbine/aot/params.py
@@ -266,9 +266,19 @@ def save(self, file_path: Union[str, Path]):
 
     def add_tensor(self, name: str, tensor: torch.Tensor):
         """Adds an named tensor to the archive."""
-        flat_array = tensor.detach().flatten().contiguous().cpu().view(torch.uint8)
-        host_array = flat_array.numpy()
-        self._index.add_buffer(name, host_array, metadata=_make_tensor_metadata(tensor))
+        metadata = _make_tensor_metadata(tensor)
+        # 0d arrays are special in both torch/numpy in different ways that makes
+        # it hard to reliably get a memory view of their contents. Since we
+        # know that 0d is always small, we just force a copy when in numpy
+        # land and that seems to get it on the happy path.
+        # See: https://github.com/iree-org/iree-turbine/issues/29
+        if len(tensor.shape) == 0:
+            flat_array = tensor.detach().cpu().numpy().copy()
+            host_array = flat_array
+        else:
+            flat_array = tensor.detach().flatten().contiguous().cpu().view(torch.uint8)
+            host_array = flat_array.numpy()
+        self._index.add_buffer(name, host_array, metadata=metadata)
 
     def add_module(self, module: nn.Module, *, prefix: str = ""):
         """Adds all parameters and persistent buffers from a module hierarchy."""


Test patch:
diff --git a/tests/aot/params_test.py b/tests/aot/params_test.py
index 5114e0ae..a1d64206 100644
--- a/tests/aot/params_test.py
+++ b/tests/aot/params_test.py
@@ -18,6 +18,7 @@
     save_module_parameters,
     ExternalTensorTrait,
     ParameterArchive,
+    ParameterArchiveBuilder,
 )
 
 
@@ -49,6 +50,32 @@ def testCreateArchive(self):
             torch.testing.assert_close(weight, m.classifier.weight)
             torch.testing.assert_close(bias, m.classifier.bias)
 
+    def testRoundtripScalarTensor(self):
+        # See: https://github.com/iree-org/iree-turbine/issues/29
+        with tempfile.TemporaryDirectory() as td:
+            file_path = Path(td) / "archive.irpa"
+            orig_scalar = torch.tensor(0.5, dtype=torch.float32)
+            builder = ParameterArchiveBuilder()
+            builder.add_tensor("scalar", orig_scalar)
+            builder.save(file_path)
+            archive = ParameterArchive(file_path, mmap=False)
+            items = dict(archive.items())
+            scalar = items["scalar"].as_tensor()
+            torch.testing.assert_close(orig_scalar, scalar)
+
+    def testRoundtripScalarUint8(self):
+        # See: https://github.com/iree-org/iree-turbine/issues/29
+        with tempfile.TemporaryDirectory() as td:
+            file_path = Path(td) / "archive.irpa"
+            orig_scalar = torch.tensor(8, dtype=torch.uint8)
+            builder = ParameterArchiveBuilder()
+            builder.add_tensor("scalar", orig_scalar)
+            builder.save(file_path)
+            archive = ParameterArchive(file_path, mmap=False)
+            items = dict(archive.items())
+            scalar = items["scalar"].as_tensor()
+            torch.testing.assert_close(orig_scalar, scalar)
+
     def testCreateArchiveWithPrefixScope(self):
         with tempfile.TemporaryDirectory() as td:
             file_path = Path(td) / "archive.irpa"