code: def gehrd_hlo(dtype, a):
      _lapack.initialize()
      a_type = ir.RankedTensorType(a.type)
      dims = a_type.shape
      assert len(dims) >= 2
      m, n = dims[-2:]
      assert m == n, (m, n)
      batch_dims = tuple(dims[:-2])
      num_bd = len(batch_dims)
      b = 1
      for d in batch_dims:
          b *= d
  
      if dtype == np.float32:
          fn = "lapack_sgehrd"
          lwork = _lapack.lapack_sgehrd_workspace(n, n, 1, n)
      elif dtype == np.float64:
          fn = "lapack_dgehrd"
          lwork = _lapack.lapack_dgehrd_workspace(n, n, 1, n)
      elif dtype == np.complex64:
          fn = "lapack_cgehrd"
          lwork = _lapack.lapack_cgehrd_workspace(n, n, 1, n)
      elif dtype == np.complex128:
          fn = "lapack_zgehrd"
          lwork = _lapack.lapack_zgehrd_workspace(n, n, 1, n)
      else:
          raise NotImplementedError(f"Unsupported dtype {dtype}")
  
      layout = (num_bd, num_bd + 1) + tuple(range(num_bd - 1, -1, -1))
      i32_type = ir.IntegerType.get_signless(32)
      out = custom_call(
          fn,
          result_types=[
              a.type,
              ir.RankedTensorType.get(
                  batch_dims + (n - 1,), a_type.element_type
              ),
              ir.RankedTensorType.get(batch_dims, i32_type),
              ir.RankedTensorType.get([lwork], a_type.element_type),
          ],
          operands=[
              hlo_s32(n),
              hlo_s32(1),
              hlo_s32(n),
              hlo_s32(n),
              hlo_s32(b),
              hlo_s32(lwork),
              a,
          ],
          operand_layouts=[[]] * 6 + [layout],
          result_layouts=[
              layout,
              (num_bd,) + tuple(range(num_bd - 1, -1, -1)),
              tuple(range(num_bd - 1, -1, -1)),
              [0],
          ],
          operand_output_aliases={6: 0},
      ).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 10 for calling a global function with example usage: gehrd_hlo(dtype, a) and returns: out