code: def syevd_hlo(
      dtype,
      a: ir.Value,
      a_shape_vals: tuple[DimensionSize, ...],
      lower=False,
  ):
      _lapack.initialize()
      a_type = ir.RankedTensorType(a.type)
      assert len(a_shape_vals) >= 2
      m, n = a_shape_vals[-2:]
      # Non-batch dimensions must be static
      assert type(m) is int and type(n) is int and m == n, a_shape_vals
  
      batch_dims_vals = a_shape_vals[:-2]
      num_bd = len(a_shape_vals) - 2
  
      i32_type = ir.IntegerType.get_signless(32)
      workspace: list[ShapeTypePair]
      if dtype == np.float32:
          fn = "lapack_ssyevd"
          eigvals_type = ir.F32Type.get()
          workspace = [
              ([_lapack.syevd_work_size(n)], a_type.element_type),
              ([_lapack.syevd_iwork_size(n)], i32_type),
          ]
      elif dtype == np.float64:
          fn = "lapack_dsyevd"
          eigvals_type = ir.F64Type.get()
          workspace = [
              ([_lapack.syevd_work_size(n)], a_type.element_type),
              ([_lapack.syevd_iwork_size(n)], i32_type),
          ]
      elif dtype == np.complex64:
          fn = "lapack_cheevd"
          eigvals_type = ir.F32Type.get()
          workspace = [
              ([_lapack.heevd_work_size(n)], a_type.element_type),
              ([_lapack.heevd_rwork_size(n)], eigvals_type),
              ([_lapack.syevd_iwork_size(n)], i32_type),
          ]
      elif dtype == np.complex128:
          fn = "lapack_zheevd"
          eigvals_type = ir.F64Type.get()
          workspace = [
              ([_lapack.heevd_work_size(n)], a_type.element_type),
              ([_lapack.heevd_rwork_size(n)], eigvals_type),
              ([_lapack.syevd_iwork_size(n)], i32_type),
          ]
      else:
          raise NotImplementedError(f"Unsupported dtype {dtype}")
  
      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)
          )
  
      scalar_layout = []
      shape_layout = [0]
      workspace_layouts = [shape_layout] * len(workspace)
      layout = (num_bd, num_bd + 1) + tuple(range(num_bd - 1, -1, -1))
  
      result_types, result_shapes = mk_result_types_and_shapes(
          [
              (a_shape_vals, a_type.element_type),
              (batch_dims_vals + (n,), eigvals_type),
              (batch_dims_vals, i32_type),
          ]
          + workspace
      )
  
      out = custom_call(
          fn,
          result_types=result_types,
          operands=[
              hlo_s32(1 if lower else 0),
              batch_size_val,
              ensure_hlo_s32(n),
              a,
          ],
          operand_layouts=[scalar_layout] * 3 + [layout],
          result_layouts=[
              layout,
              tuple(range(num_bd, -1, -1)),
              tuple(range(num_bd - 1, -1, -1)),
          ]
          + workspace_layouts,
          operand_output_aliases={3: 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 7 for calling a global function with example usage: syevd_hlo(dtype, a, a_shape_vals, lower) and returns: out