code: def _syevd_hlo(
      platform,
      gpu_solver,
      have_jacobi_solver,
      dtype,
      a,
      *,
      a_shape_vals: tuple[DimensionSize, ...],
      lower=False,
  ):
      """Symmetric (Hermitian) eigendecomposition."""
      a_type = ir.RankedTensorType(a.type)
      assert len(a_shape_vals) >= 2
      m, n = a_shape_vals[-2:]
      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(batch_dims_vals)
      layout = (num_bd, num_bd + 1) + tuple(range(num_bd - 1, -1, -1))
  
      dynamic_batch_dims = any(type(d) != int for d in batch_dims_vals)
      if dynamic_batch_dims:
          batch_int = (
              -1
          )  # Signals to the kernel that the batch is an operand.
      else:
          batch_int = math.prod(batch_dims_vals)
  
      if have_jacobi_solver and n <= 32 and not dynamic_batch_dims:
          # We cannot use syevj for dynamic shapes because the workspace size
          # depends on the batch size.
          kernel = f"{platform}solver_syevj"
          lwork, opaque = gpu_solver.build_syevj_descriptor(
              np.dtype(dtype), lower, batch_int, n
          )
      else:
          kernel = f"{platform}solver_syevd"
          lwork, opaque = gpu_solver.build_syevd_descriptor(
              np.dtype(dtype), lower, batch_int, n
          )
          assert lwork > 0
  
      if ir.ComplexType.isinstance(a_type.element_type):
          eigvals_type = ir.ComplexType(
              a_type.element_type
          ).element_type
      else:
          eigvals_type = a_type.element_type
  
      i32_type = ir.IntegerType.get_signless(32)
      operands = [a]
      operand_layouts = [layout]
      if dynamic_batch_dims:
          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)
              )
          operands.append(batch_size_val)
          operand_layouts.append(())
  
      shape_type_pairs: Sequence[ShapeTypePair] = [
          (a_shape_vals, a_type.element_type),
          (batch_dims_vals + (n,), eigvals_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(
          kernel,
          result_types=result_types,
          operands=operands,
          backend_config=opaque,
          operand_layouts=operand_layouts,
          result_layouts=[
              layout,
              tuple(range(num_bd, -1, -1)),
              tuple(range(num_bd - 1, -1, -1)),
              [0],
          ],
          operand_output_aliases={0: 0},
          result_shapes=result_shapes,
      ).results
      return out[:3]
  
file_path: jax/jaxlib/gpu_solver.py
imports: from collections.abc import Sequence
  from functools import partial
  import importlib
  import math
  import jaxlib.mlir.ir
  import jaxlib.mlir.dialects.stablehlo
  import numpy
  from .gpu_common_utils import GpuLibNotLinkedError
  from jaxlib import xla_client
  from .hlo_helpers import DimensionSize, ShapeTypePair, mk_result_types_and_shapes, custom_call, ensure_hlo_s32, hlo_s32, dense_int_array
label_desc: Use image node_id 7 for calling a global function with example usage: _syevd_hlo(platform, gpu_solver, have_jacobi_solver, dtype, a) and returns: out