code: def gees_hlo(
      dtype,
      a,
      *,
      jobvs=True,
      sort=False,
      select=None,
      a_shape_vals: tuple[DimensionSize, ...],
  ):
      _lapack.initialize()
      a_type = ir.RankedTensorType(a.type)
      etype = a_type.element_type
      assert len(a_shape_vals) >= 2
      n = a_shape_vals[-1]
      batch_dims_vals = a_shape_vals[:-2]
      num_bd = len(batch_dims_vals)
      layout = (num_bd, num_bd + 1) + tuple(range(num_bd - 1, -1, -1))
  
      if sort:
          raise NotImplementedError(
              "The sort feature of LAPACK's gees routine is not implemented."
          )
  
      jobvs = ord("V" if jobvs else "N")
      sort = ord("S" if sort else "N")
  
      if dtype == np.float32:
          fn = "lapack_sgees"
      elif dtype == np.float64:
          fn = "lapack_dgees"
      elif dtype == np.complex64:
          fn = "lapack_cgees"
      elif dtype == np.complex128:
          fn = "lapack_zgees"
      else:
          raise NotImplementedError(f"Unsupported dtype {dtype}")
  
      workspaces: list[ShapeTypePair]
      eigvals: list[ShapeTypePair]
      if not np.issubdtype(dtype, np.complexfloating):
          workspaces = [(a_shape_vals, etype)]
          workspace_layouts = [layout]
          eigvals = [(batch_dims_vals + (n,), etype)] * 2
          eigvals_layouts = [tuple(range(num_bd, -1, -1))] * 2
      else:
          workspaces = [
              (a_shape_vals, etype),
              ([n], ir.ComplexType(etype).element_type),
          ]
          workspace_layouts = [layout, [0]]
          eigvals = [(batch_dims_vals + (n,), etype)]
          eigvals_layouts = [tuple(range(num_bd, -1, -1))]
  
      i32_type = ir.IntegerType.get_signless(32)
  
      scalar_layout = []
      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 = (
          workspaces
          + eigvals
          + [
              (a_shape_vals, etype),
              (batch_dims_vals, i32_type),
              (batch_dims_vals, i32_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,
              ensure_hlo_s32(n),
              hlo_u8(jobvs),
              hlo_u8(sort),
              # TODO: figure out how to put the callable select function here
              a,
          ],
          operand_layouts=[scalar_layout] * 4 + [layout],
          result_layouts=workspace_layouts
          + eigvals_layouts
          + [
              layout,
              tuple(range(num_bd - 1, -1, -1)),
              tuple(range(num_bd - 1, -1, -1)),
          ],
          operand_output_aliases={4: 0},
          result_shapes=result_shapes,
      ).results
      if sort == ord("S"):
          return (out[0], out[3], out[4], out[5])
      else:
          return (out[0], out[3], out[5])
  
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 9 for calling a global function with example usage: gees_hlo(dtype, a) and returns: out,  out,  out,  out, out,  out,  out