code: def qr(
      a: ArrayLike, mode: str = "reduced"
  ) -> Array | tuple[Array, Array]:
      check_arraylike("jnp.linalg.qr", a)
      (a,) = promote_dtypes_inexact(jnp.asarray(a))
      if mode == "raw":
          a, taus = lax_linalg.geqrf(a)
          return a.mT, taus
      if mode in ("reduced", "r", "full"):
          full_matrices = False
      elif mode == "complete":
          full_matrices = True
      else:
          raise ValueError(
              f"Unsupported QR decomposition mode '{mode}'"
          )
      q, r = lax_linalg.qr(a, full_matrices=full_matrices)
      if mode == "r":
          return r
      return q, r
  
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 31 for calling a global function with example usage: qr(a, mode) and returns: q,  r, a,  taus, r