jax.numpy.linalg.lstsq throws IndexError with zero rows
### Description

[jax.numpy.linalg.lstsq](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.linalg.lstsq.html) throws an IndexError when the number of rows is zero:

```python3
import numpy as np
import jax.numpy as jnp

A = np.zeros([0, 3])
b = np.zeros([0])
x, residuals, rank, s = np.linalg.lstsq(A, b, None)
print(x) # [0. 0. 0.]

A = jnp.zeros([0, 3])
b = jnp.zeros([0])
x, residuals, rank, s = jnp.linalg.lstsq(A, b, None)
# IndexError: index is out of bounds for axis 0 with size 0
```

### What jax/jaxlib version are you using?

jax 0.3.21, jaxlib 0.3.20

### Which accelerator(s) are you using?

CPU

### Additional system info

Python 3.10.7, macOS 11.7

### NVIDIA GPU info

_No response_
diff --git a/jax/_src/numpy/linalg.py b/jax/_src/numpy/linalg.py
index 49cbbeb51..0c5f07df4 100644
--- a/jax/_src/numpy/linalg.py
+++ b/jax/_src/numpy/linalg.py
@@ -601,17 +601,22 @@ def _lstsq(a: ArrayLike, b: ArrayLike, rcond: Optional[float], *,
       f"{b.ndim}-dimensional array given. Array must be one or two-dimensional")
   m, n = a.shape
   dtype = a.dtype
-  if rcond is None:
-    rcond = jnp.finfo(dtype).eps * max(n, m)
+  if a.size == 0:
+    s = jnp.empty(0, dtype=a.dtype)
+    rank = jnp.array(0, dtype=int)
+    x = jnp.empty((n, *b.shape[1:]), dtype=a.dtype)
   else:
-    rcond = jnp.where(rcond < 0, jnp.finfo(dtype).eps, rcond)
-  u, s, vt = svd(a, full_matrices=False)
-  mask = s >= jnp.array(rcond, dtype=s.dtype) * s[0]
-  rank = mask.sum()
-  safe_s = jnp.where(mask, s, 1).astype(a.dtype)
-  s_inv = jnp.where(mask, 1 / safe_s, 0)[:, jnp.newaxis]
-  uTb = jnp.matmul(u.conj().T, b, precision=lax.Precision.HIGHEST)
-  x = jnp.matmul(vt.conj().T, s_inv * uTb, precision=lax.Precision.HIGHEST)
+    if rcond is None:
+      rcond = jnp.finfo(dtype).eps * max(n, m)
+    else:
+      rcond = jnp.where(rcond < 0, jnp.finfo(dtype).eps, rcond)
+    u, s, vt = svd(a, full_matrices=False)
+    mask = s >= jnp.array(rcond, dtype=s.dtype) * s[0]
+    rank = mask.sum()
+    safe_s = jnp.where(mask, s, 1).astype(a.dtype)
+    s_inv = jnp.where(mask, 1 / safe_s, 0)[:, jnp.newaxis]
+    uTb = jnp.matmul(u.conj().T, b, precision=lax.Precision.HIGHEST)
+    x = jnp.matmul(vt.conj().T, s_inv * uTb, precision=lax.Precision.HIGHEST)
   # Numpy returns empty residuals in some cases. To allow compilation, we
   # default to returning full residuals in all cases.
   if numpy_resid and (rank < n or m <= n):
diff --git a/tests/linalg_test.py b/tests/linalg_test.py
index 0169380c0..7767f4619 100644
--- a/tests/linalg_test.py
+++ b/tests/linalg_test.py
@@ -890,6 +890,9 @@ class NumpyLinalgTest(jtu.JaxTestCase):
           ((4, 6), (4,)),
           ((6, 6), (6, 1)),
           ((8, 6), (8, 4)),
+          ((0, 3), (0,)),
+          ((3, 0), (3,)),
+          ((3, 1), (3, 0)),
       ]
     ],
     rcond=[-1, None, 0.5],
['tests/linalg_test.py', 'jax/_src/numpy/linalg.py']