Merge pull request #18833 from jakevdp:linalg-shapes

PiperOrigin-RevId: 588458694
jax-ml · Dec 6, 2023 · 9e3a8fa · 9e3a8fa
2 parents 4bdcb11 + 45905fa
commit 9e3a8fa
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git a/jax/_src/numpy/linalg.py b/jax/_src/numpy/linalg.py
@@ -606,6 +606,8 @@ def qr(a: ArrayLike, mode: str = "reduced") -> Union[Array, tuple[Array, Array]]
 def solve(a: ArrayLike, b: ArrayLike) -> Array:
   check_arraylike("jnp.linalg.solve", a, b)
   a, b = promote_dtypes_inexact(jnp.asarray(a), jnp.asarray(b))
+  if a.ndim >= 2 and b.ndim > a.ndim:
+    a = lax.expand_dims(a, tuple(range(b.ndim - a.ndim)))
   return lax_linalg._solve(a, b)
 
 

diff --git a/tests/linalg_test.py b/tests/linalg_test.py
@@ -897,6 +897,7 @@ def tensor_maker():
       ((4, 4), (4,)),
       ((8, 8), (8, 4)),
       ((1, 2, 2), (3, 2)),
+      ((2, 2), (3, 2, 2)),
       ((2, 1, 3, 3), (1, 4, 3, 4)),
       ((1, 0, 0), (1, 0, 2)),
      ]