Add torch wrappers for ones(), zeros(), and empty()

This allows shape to be passed as a keyword argument.

Fixes #22.

Author: asmeurer

array_api_compat/torch/_aliases.py

@@ -519,6 +519,28 @@ def full(shape: Union[int, Tuple[int, ...]],
 
     return torch.full(shape, fill_value, dtype=dtype, device=device, **kwargs)
 
+# ones, zeros, and empty do not accept shape as a keyword argument
+def ones(shape: Union[int, Tuple[int, ...]],
+         *,
+         dtype: Optional[Dtype] = None,
+         device: Optional[Device] = None,
+         **kwargs) -> array:
+    return torch.ones(shape, dtype=dtype, device=device, **kwargs)
+
+def zeros(shape: Union[int, Tuple[int, ...]],
+         *,
+         dtype: Optional[Dtype] = None,
+         device: Optional[Device] = None,
+         **kwargs) -> array:
+    return torch.zeros(shape, dtype=dtype, device=device, **kwargs)
+
+def empty(shape: Union[int, Tuple[int, ...]],
+         *,
+         dtype: Optional[Dtype] = None,
+         device: Optional[Device] = None,
+         **kwargs) -> array:
+    return torch.empty(shape, dtype=dtype, device=device, **kwargs)
+
 # Functions that aren't in torch https://github.com/pytorch/pytorch/issues/58742
 def expand_dims(x: array, /, *, axis: int = 0) -> array:
     return torch.unsqueeze(x, axis)
@@ -585,7 +607,7 @@ def tensordot(x1: array, x2: array, /, *, axes: Union[int, Tuple[Sequence[int],
            'logaddexp', 'multiply', 'not_equal', 'pow', 'remainder',
            'subtract', 'max', 'min', 'sort', 'prod', 'sum', 'any', 'all',
            'mean', 'std', 'var', 'concat', 'squeeze', 'flip', 'roll',
-           'nonzero', 'where', 'arange', 'eye', 'linspace', 'full',
-           'expand_dims', 'astype', 'broadcast_arrays', 'unique_all',
-           'unique_counts', 'unique_inverse', 'unique_values',
+           'nonzero', 'where', 'arange', 'eye', 'linspace', 'full', 'ones',
+           'zeros', 'empty', 'expand_dims', 'astype', 'broadcast_arrays',
+           'unique_all', 'unique_counts', 'unique_inverse', 'unique_values',
            'matmul', 'matrix_transpose', 'vecdot', 'tensordot']