Merge pull request jax-ml#25338 from carlosgmartin:fix_numpy_linalg_matrix_norm_ord_type_annotation

PiperOrigin-RevId: 704245037

Author: Google-ML-Automation

jax/_src/numpy/linalg.py

@@ -1517,7 +1517,7 @@ def outer(x1: ArrayLike, x2: ArrayLike, /) -> Array:
 
 
 @export
-def matrix_norm(x: ArrayLike, /, *, keepdims: bool = False, ord: str = 'fro') -> Array:
+def matrix_norm(x: ArrayLike, /, *, keepdims: bool = False, ord: str | int = 'fro') -> Array:
   """Compute the norm of a matrix or stack of matrices.
 
   JAX implementation of :func:`numpy.linalg.matrix_norm`