Match PyTorch ones API

iris/iris.py

@@ -362,41 +362,61 @@ def randn(
     def ones(self, *size, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False):
         """
         Returns a tensor filled with the scalar value 1, with the shape defined by the variable argument size.
+        The tensor is allocated on the Iris symmetric heap.
 
         Args:
-            *size (int...): a sequence of integers defining the shape of the output tensor. Can be a variable number of arguments or a collection like a list or tuple.
+            *size (int...): a sequence of integers defining the shape of the output tensor.
+                Can be a variable number of arguments or a collection like a list or tuple.
 
         Keyword Arguments:
             out (Tensor, optional): the output tensor.
-            dtype (torch.dtype, optional): the desired data type of returned tensor. Default: if None, uses a global default (see torch.set_default_dtype()).
-            layout (torch.layout, optional): the desired layout of returned Tensor. Default: torch.strided.
-            device (torch.device, optional): the desired device of returned tensor. Default: if None, uses the current device for the default tensor type.
-            requires_grad (bool, optional): If autograd should record operations on the returned tensor. Default: False.
+            dtype (torch.dtype, optional): the desired data type of returned tensor.
+                Default: if None, uses a global default (see torch.set_default_dtype()).
+            layout (torch.layout, optional): the desired layout of returned Tensor.
+                Default: torch.strided. Note: Iris tensors always use `torch.strided` regardless of this parameter.
+            device (torch.device, optional): the desired device of returned tensor.
+                Default: if None, uses the current device for the default tensor type.
+            requires_grad (bool, optional): If autograd should record operations on the returned tensor.
+                Default: False.
         """
         self.debug(f"ones: size = {size}, dtype = {dtype}, device = {device}, requires_grad = {requires_grad}")
 
-        # Validate device
-        self.__throw_if_invalid_device(device)
-
-        # Handle the case where size is provided as a single tuple/list
-        if len(size) == 1 and isinstance(size[0], (tuple, list)):
-            size = size[0]
-
         # Use global default dtype if None is provided
         if dtype is None:
             dtype = torch.get_default_dtype()
+
+        # Use current device if none specified
+        if device is None:
+            device = self.device
+
+        # Validate device compatibility with Iris
+        self.__throw_if_invalid_device(device)
+
+        # Parse size and calculate number of elements
         size, num_elements = self.parse_size(size)
 
         # If out is provided, use it; otherwise allocate new tensor
         if out is not None:
             self.__throw_if_invalid_output_tensor(out, num_elements, dtype)
-            tensor = out
+            # Fill with ones
+            out.fill_(1)
+            # Create a reshaped view of the out tensor
+            tensor = out.view(size)
         else:
             tensor = self.allocate(num_elements=num_elements, dtype=dtype)
-        tensor.fill_(1)
+            # Fill with ones
+            tensor.fill_(1)
+            # Reshape to the desired size
+            tensor = tensor.reshape(size)
+
+        # Apply the requested layout
+        tensor = self.__apply_layout(tensor, layout)
+
+        # Set requires_grad if specified
         if requires_grad:
             tensor.requires_grad_()
-        return tensor.reshape(size)
+
+        return tensor
 
     def full(self, size, fill_value, dtype=torch.int):
         self.debug(f"full: size = {size}, fill_value = {fill_value}, dtype = {dtype}")