Move linalg aliases to torch/_aliases

Author: adonath

array_api_compat/torch/_aliases.py

@@ -1,22 +1,24 @@
 from __future__ import annotations
 
+from builtins import all as builtin_all
+from builtins import any as builtin_any
 from functools import wraps
-from builtins import all as builtin_all, any as builtin_any
-
-from ..common._aliases import (UniqueAllResult, UniqueCountsResult,
-                               UniqueInverseResult,
-                               matrix_transpose as _aliases_matrix_transpose,
-                               vecdot as _aliases_vecdot)
-from .._internal import get_xp
+from typing import TYPE_CHECKING
 
 import torch
 
-from typing import TYPE_CHECKING
+from .._internal import get_xp
+from ..common._aliases import UniqueAllResult, UniqueCountsResult, UniqueInverseResult
+from ..common._aliases import matrix_transpose as _aliases_matrix_transpose
+from ..common._aliases import vecdot as _aliases_vecdot
+
 if TYPE_CHECKING:
     from typing import List, Optional, Sequence, Tuple, Union
-    from ..common._typing import Device
+
     from torch import dtype as Dtype
 
+    from ..common._typing import Device
+
     array = torch.Tensor
 
 _int_dtypes = {
@@ -693,15 +695,42 @@ def take(x: array, indices: array, /, *, axis: Optional[int] = None, **kwargs) -
         axis = 0
     return torch.index_select(x, axis, indices, **kwargs)
 
-__all__ = ['result_type', 'can_cast', 'permute_dims', 'bitwise_invert', 'newaxis',
-           'add', 'atan2', 'bitwise_and', 'bitwise_left_shift', 'bitwise_or',
-           'bitwise_right_shift', 'bitwise_xor', 'divide', 'equal',
-           'floor_divide', 'greater', 'greater_equal', 'less', 'less_equal',
-           'logaddexp', 'multiply', 'not_equal', 'pow', 'remainder',
-           'subtract', 'max', 'min', 'sort', 'prod', 'sum', 'any', 'all',
-           'mean', 'std', 'var', 'concat', 'squeeze', 'broadcast_to', 'flip', 'roll',
-           'nonzero', 'where', 'reshape', 'arange', 'eye', 'linspace', 'full',
-           'ones', 'zeros', 'empty', 'tril', 'triu', 'expand_dims', 'astype',
-           'broadcast_arrays', 'unique_all', 'unique_counts',
-           'unique_inverse', 'unique_values', 'matmul', 'matrix_transpose',
-           'vecdot', 'tensordot', 'isdtype', 'take']
+
+
+# Note: torch.linalg.cross does not default to axis=-1 (it defaults to the
+# first axis with size 3), see https://github.com/pytorch/pytorch/issues/58743
+def cross(x1: array, x2: array, /, *, axis: int = -1) -> array:
+    x1, x2 = _fix_promotion(x1, x2, only_scalar=False)
+    return torch.linalg.cross(x1, x2, dim=axis)
+
+def vecdot_linalg(x1: array, x2: array, /, *, axis: int = -1, **kwargs) -> array:
+    from ._aliases import isdtype
+
+    x1, x2 = _fix_promotion(x1, x2, only_scalar=False)
+
+    # torch.linalg.vecdot doesn't support integer dtypes
+    if isdtype(x1.dtype, 'integral') or isdtype(x2.dtype, 'integral'):
+        if kwargs:
+            raise RuntimeError("vecdot kwargs not supported for integral dtypes")
+        ndim = max(x1.ndim, x2.ndim)
+        x1_shape = (1,)*(ndim - x1.ndim) + tuple(x1.shape)
+        x2_shape = (1,)*(ndim - x2.ndim) + tuple(x2.shape)
+        if x1_shape[axis] != x2_shape[axis]:
+            raise ValueError("x1 and x2 must have the same size along the given axis")
+
+        x1_, x2_ = torch.broadcast_tensors(x1, x2)
+        x1_ = torch.moveaxis(x1_, axis, -1)
+        x2_ = torch.moveaxis(x2_, axis, -1)
+
+        res = x1_[..., None, :] @ x2_[..., None]
+        return res[..., 0, 0]
+    return torch.linalg.vecdot(x1, x2, dim=axis, **kwargs)
+
+def solve(x1: array, x2: array, /, **kwargs) -> array:
+    x1, x2 = _fix_promotion(x1, x2, only_scalar=False)
+    return torch.linalg.solve(x1, x2, **kwargs)
+
+# torch.trace doesn't support the offset argument and doesn't support stacking
+def trace(x: array, /, *, offset: int = 0, dtype: Optional[Dtype] = None) -> array:
+    # Use our wrapped sum to make sure it does upcasting correctly
+    return sum(torch.diagonal(x, offset=offset, dim1=-2, dim2=-1), axis=-1, dtype=dtype)