Use Literal for ord typehint

Author: jessegrabowski

pytensor/tensor/nlinalg.py

@@ -1,6 +1,6 @@
 import warnings
 from functools import partial
-from typing import Callable, Optional, Union
+from typing import Callable, Literal, Optional, Union
 
 import numpy as np
 from numpy.core.numeric import normalize_axis_tuple  # type: ignore
@@ -723,9 +723,12 @@ def _multi_svd_norm(
     return result
 
 
+VALID_ORD = Literal["fro", "f", "nuc", "inf", "-inf", 0, 1, -1, 2, -2]
+
+
 def norm(
     x: ptb.TensorVariable,
-    ord: Optional[Union[float, str]] = None,
+    ord: Optional[Union[float, VALID_ORD]] = None,
     axis: Optional[Union[int, tuple[int, ...]]] = None,
     keepdims: bool = False,
 ):
@@ -736,16 +739,22 @@ def norm(
     ----------
     x: TensorVariable
         Tensor to take norm of.
-    ord: float or str, optional
+
+    ord: float, str or int, optional
         Order of norm. If `ord` is a str, it must be one of the following:
             - 'fro' or 'f' : Frobenius norm
             - 'nuc' : nuclear norm
             - 'inf' : Infinity norm
             - '-inf' : Negative infinity norm
-        Otherwise `ord` must be a float. Default is the Frobenius (L2) norm.
+        If an integer, order can be one of -2, -1, 0, 1, or 2.
+        Otherwise `ord` must be a float.
+
+        Default is the Frobenius (L2) norm.
+
     axis: tuple of int, optional
         Axes over which to compute the norm. If None, norm of entire matrix (or vector) is computed. Row or column
         norms can be computed by passing a single integer; this will treat a matrix like a batch of vectors.
+
     keepdims: bool
         If True, dummy axes will be inserted into the output so that norm.dnim == x.dnim. Default is False.