TYP: reject bool in the ord params of vector_norm and matrix_norm

Author: jorenham

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: JustInt | 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/torch/linalg.py

@@ -16,6 +16,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
@@ -84,8 +85,8 @@ def vector_norm(
     *,
     axis: Optional[Union[int, Tuple[int, ...]]] = None,
     keepdims: bool = False,
-    # float stands for inf | -inf, which are not valid for Literal
-    ord: Union[int, float] = 2,
+    # JustFloat stands for inf | -inf, which are not valid for Literal
+    ord: JustInt | JustFloat = 2,
     **kwargs,
 ) -> Array:
     # torch.vector_norm incorrectly treats axis=() the same as axis=None
@@ -115,3 +116,6 @@ def vector_norm(
 _all_ignore = ['torch_linalg', 'sum']
 
 del linalg_all
+
+def __dir__() -> list[str]:
+    return __all__