code: def geev_hlo(
      dtype,
      input,
      *,
      input_shape_vals: tuple[
          DimensionSize, ...
      ],  # input.shape as ir.Values
      jobvl=True,
      jobvr=True,
  ):
      # input_shape_vals are used for when input has dynamic shapes.
      _lapack.initialize()
      input_shape = ir.RankedTensorType(input.type).shape
      assert len(input_shape) >= 2
      n = input_shape_vals[-1]
      batch_dims_vals = input_shape_vals[:-2]
      num_bd = len(batch_dims_vals)
  
      layout = (num_bd, num_bd + 1) + tuple(range(num_bd - 1, -1, -1))
  
      jobvl_c = ord("V" if jobvl else "N")
      jobvr_c = ord("V" if jobvr else "N")
  
      i32_type = ir.IntegerType.get_signless(32)
      f32_type = ir.F32Type.get()
      f64_type = ir.F64Type.get()
      c64_type = ir.ComplexType.get(ir.F32Type.get())
      c128_type = ir.ComplexType.get(ir.F64Type.get())
  
      workspaces: list[ShapeTypePair]
      eigvals: list[ShapeTypePair]
      if dtype == np.float32:
          fn = "lapack_sgeev"
          real = True
          eigvecs_type = c64_type
          workspaces = [([n, n], f32_type)] * 3
          workspace_layouts = [[0, 1]] * 3
          eigvals = [(batch_dims_vals + (n,), f32_type)] * 2
          eigvals_layouts = [tuple(range(num_bd, -1, -1))] * 2
      elif dtype == np.float64:
          fn = "lapack_dgeev"
          real = True
          eigvecs_type = c128_type
          workspaces = [([n, n], f64_type)] * 3
          workspace_layouts = [[0, 1]] * 3
          eigvals = [(batch_dims_vals + (n,), f64_type)] * 2
          eigvals_layouts = [tuple(range(num_bd, -1, -1))] * 2
      elif dtype == np.complex64:
          fn = "lapack_cgeev"
          real = False
          eigvecs_type = c64_type
          workspaces = [([n, n], c64_type), ([hlo_add(n, n)], f32_type)]
          workspace_layouts = [[0, 1], [0]]
          eigvals = [(batch_dims_vals + (n,), c64_type)]
          eigvals_layouts = [tuple(range(num_bd, -1, -1))]
      elif dtype == np.complex128:
          fn = "lapack_zgeev"
          real = False
          eigvecs_type = c128_type
          workspaces = [
              ([n, n], c128_type),
              ([hlo_add(n, n)], f64_type),
          ]
          workspace_layouts = [[0, 1], [0]]
          eigvals = [(batch_dims_vals + (n,), c128_type)]
          eigvals_layouts = [tuple(range(num_bd, -1, -1))]
      else:
          raise NotImplementedError(f"Unsupported dtype {dtype}")
  
      scalar_layout = []
      info_layout = tuple(range(num_bd - 1, -1, -1))
  
      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] = (
          workspaces
          + eigvals
          + [
              (input_shape_vals, eigvecs_type),
              (input_shape_vals, eigvecs_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(jobvl_c),
              hlo_u8(jobvr_c),
              input,
          ],
          operand_layouts=[scalar_layout] * 4 + [layout],
          result_layouts=(
              workspace_layouts
              + eigvals_layouts
              + [layout] * 2
              + [info_layout]
          ),
          result_shapes=result_shapes,
      ).results
      if real:
          return (hlo.complex(out[3], out[4]), out[5], out[6], out[7])
      else:
          return out[2:6]
  
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 8 for calling a global function with example usage: geev_hlo(dtype, input) and returns: hlo,  out,  out,  out, out