code: def _geqrf_hlo(platform, gpu_solver, dtype, a):
      """QR decomposition."""
      a_type = ir.RankedTensorType(a.type)
      dims = a_type.shape
      assert len(dims) >= 2
      m, n = dims[-2:]
      batch_dims = tuple(dims[:-2])
      num_bd = len(batch_dims)
      batch = math.prod(batch_dims)
  
      lwork, opaque = gpu_solver.build_geqrf_descriptor(
          np.dtype(dtype), batch, m, n
      )
  
      layout = (num_bd, num_bd + 1) + tuple(range(num_bd - 1, -1, -1))
      i32_type = ir.IntegerType.get_signless(32)
      out = custom_call(
          f"{platform}solver_geqrf",
          result_types=[
              a.type,
              ir.RankedTensorType.get(
                  batch_dims + (min(m, n),), a_type.element_type
              ),
              ir.RankedTensorType.get(batch_dims, i32_type),
              ir.RankedTensorType.get([lwork], a_type.element_type),
          ],
          operands=[a],
          backend_config=opaque,
          operand_layouts=[layout],
          result_layouts=[
              layout,
              tuple(range(num_bd, -1, -1)),
              tuple(range(num_bd - 1, -1, -1)),
              [0],
          ],
          operand_output_aliases={0: 0},
      ).results
      return out[:3]
  
file_path: jax/jaxlib/gpu_solver.py
imports: from collections.abc import Sequence
  from functools import partial
  import importlib
  import math
  import jaxlib.mlir.ir
  import jaxlib.mlir.dialects.stablehlo
  import numpy
  from .gpu_common_utils import GpuLibNotLinkedError
  from jaxlib import xla_client
  from .hlo_helpers import DimensionSize, ShapeTypePair, mk_result_types_and_shapes, custom_call, ensure_hlo_s32, hlo_s32, dense_int_array
label_desc: Use image node_id 3 for calling a global function with example usage: _geqrf_hlo(platform, gpu_solver, dtype, a) and returns: out