code: def geqrf_hlo(
      dtype, a: ir.Value, *, a_shape_vals: tuple[DimensionSize, ...]
  ):
      _lapack.initialize()
      a_type = ir.RankedTensorType(a.type)
      assert len(a_shape_vals) >= 2
      m, n = a_shape_vals[-2:]
      assert type(m) is int
      assert type(n) is int
  
      batch_dims_vals = a_shape_vals[:-2]
      num_bd = len(batch_dims_vals)
  
      if dtype == np.float32:
          fn = "lapack_sgeqrf"
          lwork = _lapack.lapack_sgeqrf_workspace(m, n)
      elif dtype == np.float64:
          fn = "lapack_dgeqrf"
          lwork = _lapack.lapack_dgeqrf_workspace(m, n)
      elif dtype == np.complex64:
          fn = "lapack_cgeqrf"
          lwork = _lapack.lapack_cgeqrf_workspace(m, n)
      elif dtype == np.complex128:
          fn = "lapack_zgeqrf"
          lwork = _lapack.lapack_zgeqrf_workspace(m, n)
      else:
          raise NotImplementedError(f"Unsupported dtype {dtype}")
  
      scalar_layout = []
      layout = (num_bd, num_bd + 1) + tuple(range(num_bd - 1, -1, -1))
      i32_type = ir.IntegerType.get_signless(32)
  
      batch_size_val = hlo_s32(1)
      for b_v in batch_dims_vals:
          batch_size_val = hlo.multiply(
              batch_size_val, ensure_hlo_s32(b_v)
          )
      shape_type_pairs: Sequence[ShapeTypePair] = [
          (a_shape_vals, a_type.element_type),
          (batch_dims_vals + (min(m, n),), a_type.element_type),
          (batch_dims_vals, i32_type),
          ([lwork], a_type.element_type),
      ]
      result_types, result_shapes = mk_result_types_and_shapes(
          shape_type_pairs
      )
      out = custom_call(
          fn,
          result_types=result_types,
          operands=[
              batch_size_val,
              hlo_s32(m),
              hlo_s32(n),
              hlo_s32(lwork),
              a,
          ],
          operand_layouts=[scalar_layout] * 4 + [layout],
          result_layouts=[
              layout,
              tuple(range(num_bd, -1, -1)),
              tuple(range(num_bd - 1, -1, -1)),
              [0],
          ],
          operand_output_aliases={4: 0},
          result_shapes=result_shapes,
      ).results
      return out[:3]
  
file_path: jax/jaxlib/lapack.py
imports: from collections.abc import Sequence
  import numpy
  import jaxlib.mlir.ir
  import jaxlib.mlir.dialects.stablehlo
  from jaxlib import xla_client
  from .hlo_helpers import custom_call, hlo_u8, hlo_s32, ensure_hlo_s32, hlo_add, hlo_min, DimensionSize, ShapeTypePair, mk_result_types_and_shapes
  from .cpu import _lapack
label_desc: Use image node_id 3 for calling a global function with example usage: geqrf_hlo(dtype, a) and returns: out