code: def _lstsq(matrix, rhs, l2_regularizer=0.0, fast=True, name=None):
      """JAX/NumPy variant of tf.linalg.lstsq."""
      del l2_regularizer, name
      if fast:
          raise ValueError("`fast=True` is not supported")
      if matrix.ndim > 2:
          if JAX_MODE:
              import jax  # pylint: disable=g-import-not-at-top
  
              return jax.vmap(functools.partial(_lstsq, fast=False))(
                  matrix, rhs
              )
          return np.array(
              [
                  _lstsq(mat, r, fast=False)
                  for mat, r in zip(matrix, rhs)
              ]
          )
      rcond = None
      if JAX_MODE and matrix.dtype == np.float32:
          rcond = 0.0  # https://github.com/google/jax/issues/15591
      return np.linalg.lstsq(matrix, rhs, rcond=rcond)[0]
  
imports: import collections
  import functools
  import numpy
  from tensorflow_probability.python.internal.backend.numpy import _utils
  from tensorflow_probability.python.internal.backend.numpy import ops
label_desc: Use image node_id 12 for calling a global function with example usage: _lstsq(matrix, rhs, l2_regularizer, fast, name) and returns: np, np, jax