Merge branch 'main' into typ_v4

Author: crusaderky

array_api_compat/common/_linalg.py

@@ -12,7 +12,7 @@
 
 from .._internal import get_xp
 from ._aliases import isdtype, matmul, matrix_transpose, tensordot, vecdot
-from ._typing import Array, DType, Namespace
+from ._typing import Array, DType, JustFloat, JustInt, Namespace
 
 
 # These are in the main NumPy namespace but not in numpy.linalg
@@ -139,7 +139,7 @@ def matrix_norm(
     xp: Namespace,
     *,
     keepdims: bool = False,
-    ord: float | Literal["fro", "nuc"] | None = "fro",
+    ord: Literal[1, 2, -1, -2] | JustFloat | Literal["fro", "nuc"] | None = "fro",
 ) -> Array:
     return xp.linalg.norm(x, axis=(-2, -1), keepdims=keepdims, ord=ord)
 
@@ -155,7 +155,7 @@ def vector_norm(
     *,
     axis: int | tuple[int, ...] | None = None,
     keepdims: bool = False,
-    ord: float = 2,
+    ord: JustInt | JustFloat = 2,
 ) -> Array:
     # xp.linalg.norm tries to do a matrix norm whenever axis is a 2-tuple or
     # when axis=None and the input is 2-D, so to force a vector norm, we make

array_api_compat/common/_typing.py

@@ -1,7 +1,15 @@
 from __future__ import annotations
 
 from types import ModuleType as Namespace
-from typing import TYPE_CHECKING, Literal, Protocol, TypeAlias, TypedDict, TypeVar
+from typing import (
+    TYPE_CHECKING,
+    Literal,
+    Protocol,
+    TypeAlias,
+    TypedDict,
+    TypeVar,
+    final,
+)
 
 if TYPE_CHECKING:
     from _typeshed import Incomplete
@@ -20,6 +28,37 @@
 _T_co = TypeVar("_T_co", covariant=True)
 
 
+# These "Just" types are equivalent to the `Just` type from the `optype` library,
+# apart from them not being `@runtime_checkable`.
+# - docs: https://github.com/jorenham/optype/blob/master/README.md#just
+# - code: https://github.com/jorenham/optype/blob/master/optype/_core/_just.py
+@final
+class JustInt(Protocol):
+    @property
+    def __class__(self, /) -> type[int]: ...
+    @__class__.setter
+    def __class__(self, value: type[int], /) -> None: ...  # pyright: ignore[reportIncompatibleMethodOverride]
+
+
+@final
+class JustFloat(Protocol):
+    @property
+    def __class__(self, /) -> type[float]: ...
+    @__class__.setter
+    def __class__(self, value: type[float], /) -> None: ...  # pyright: ignore[reportIncompatibleMethodOverride]
+
+
+@final
+class JustComplex(Protocol):
+    @property
+    def __class__(self, /) -> type[complex]: ...
+    @__class__.setter
+    def __class__(self, value: type[complex], /) -> None: ...  # pyright: ignore[reportIncompatibleMethodOverride]
+
+
+#
+
+
 class NestedSequence(Protocol[_T_co]):
     def __getitem__(self, key: int, /) -> _T_co | NestedSequence[_T_co]: ...
     def __len__(self, /) -> int: ...
@@ -78,6 +117,9 @@ def shape(self, /) -> tuple[_T_co, ...]: ...
     "Device",
     "HasShape",
     "Namespace",
+    "JustInt",
+    "JustFloat",
+    "JustComplex",
     "NestedSequence",
     "SupportsArrayNamespace",
     "SupportsBufferProtocol",

array_api_compat/torch/linalg.py

@@ -14,6 +14,7 @@
 # These functions are in both the main and linalg namespaces
 from ._aliases import matmul, matrix_transpose, tensordot
 from ._typing import Array, DType
+from ..common._typing import JustInt, JustFloat
 
 # 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
@@ -82,8 +83,8 @@ def vector_norm(
     *,
     axis: int | tuple[int, ...] | None = None,
     keepdims: bool = False,
-    # float stands for inf | -inf, which are not valid for Literal
-    ord: float = 2,
+    # JustFloat stands for inf | -inf, which are not valid for Literal
+    ord: JustInt | JustFloat = 2,
     **kwargs: object,
 ) -> Array:
     # torch.vector_norm incorrectly treats axis=() the same as axis=None
@@ -113,3 +114,6 @@ def vector_norm(
 _all_ignore = ['torch_linalg', 'sum']
 
 del linalg_all
+
+def __dir__() -> list[str]:
+    return __all__