code: def _lstsq(
      a: ArrayLike,
      b: ArrayLike,
      rcond: float | None,
      *,
      numpy_resid: bool = False,
  ) -> tuple[Array, Array, Array, Array]:
      # TODO: add lstsq to lax_linalg and implement this function via those wrappers.
      # TODO: add custom jvp rule for more robust lstsq differentiation
      a, b = promote_dtypes_inexact(a, b)
      if a.shape[0] != b.shape[0]:
          raise ValueError(
              "Leading dimensions of input arrays must match"
          )
      b_orig_ndim = b.ndim
      if b_orig_ndim == 1:
          b = b[:, None]
      if a.ndim != 2:
          raise TypeError(
              f"{a.ndim}-dimensional array given. Array must be two-dimensional"
          )
      if b.ndim != 2:
          raise TypeError(
              f"{b.ndim}-dimensional array given. Array must be one or two-dimensional"
          )
      m, n = a.shape
      dtype = a.dtype
      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:
          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):
          resid = jnp.asarray([])
      else:
          b_estimate = jnp.matmul(a, x, precision=lax.Precision.HIGHEST)
          resid = norm(b - b_estimate, axis=0) ** 2
      if b_orig_ndim == 1:
          x = x.ravel()
      return x, resid, rank, s
  
imports: from __future__ import annotations
  from collections.abc import Sequence
  from functools import partial
  import numpy
  import textwrap
  import operator
  from typing import Literal, cast, overload
  import jax
  from jax import jit, custom_jvp
  from jax import lax
  from jax._src.lax import lax
  from jax._src.lax import linalg
  from jax._src.numpy import lax_numpy
  from jax._src.numpy import reductions, ufuncs
  from jax._src.numpy.util import _wraps, promote_dtypes_inexact, check_arraylike
  from jax._src.util import canonicalize_axis
  from jax._src.typing import ArrayLike, Array
label_desc: Use image node_id 33 for calling a global function with example usage: _lstsq(a, b, rcond) and returns: x,  resid,  rank,  s