Merge pull request #25338 from carlosgmartin:fix_numpy_linalg_matrix_…

…norm_ord_type_annotation PiperOrigin-RevId: 704245037
jax-ml · Dec 9, 2024 · 5a1c4c5 · 5a1c4c5
2 parents d474fed + efa35ea
commit 5a1c4c5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/jax/_src/numpy/linalg.py b/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`