code: def gesdd_hlo(
      dtype,
      a: ir.Value,
      *,
      full_matrices=True,
      compute_uv=True,
      a_shape_vals: tuple[DimensionSize, ...],
  ):
      _lapack.initialize()
      a_type = ir.RankedTensorType(a.type)
      assert len(a_shape_vals) >= 2
      m, n = a_shape_vals[-2:]
      assert type(m) is int
      assert type(n) is int
      batch_dims_vals = a_shape_vals[:-2]
      num_bd = len(batch_dims_vals)
      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)
          )
  
      i32_type = ir.IntegerType.get_signless(32)
      workspace: list[ShapeTypePair]
      if dtype == np.float32:
          fn = "lapack_sgesdd"
          singular_vals_type = ir.F32Type.get()
          lwork = _lapack.sgesdd_work_size(
              m, n, compute_uv, full_matrices
          )
          workspace = [
              ([_lapack.gesdd_iwork_size(m, n)], i32_type),
              ([lwork], a_type.element_type),
          ]
          workspace_layouts = [[0], [0]]
      elif dtype == np.float64:
          fn = "lapack_dgesdd"
          singular_vals_type = ir.F64Type.get()
          lwork = _lapack.dgesdd_work_size(
              m, n, compute_uv, full_matrices
          )
          workspace = [
              ([_lapack.gesdd_iwork_size(m, n)], i32_type),
              ([lwork], a_type.element_type),
          ]
          workspace_layouts = [[0], [0]]
      elif dtype == np.complex64:
          fn = "lapack_cgesdd"
          singular_vals_type = ir.F32Type.get()
          lwork = _lapack.cgesdd_work_size(
              m, n, compute_uv, full_matrices
          )
          workspace = [
              ([_lapack.gesdd_iwork_size(m, n)], i32_type),
              (
                  [_lapack.cgesdd_rwork_size(m, n, int(compute_uv))],
                  ir.F32Type.get(),
              ),
              ([lwork], a_type.element_type),
          ]
          workspace_layouts = [[0], [0], [0]]
      elif dtype == np.complex128:
          fn = "lapack_zgesdd"
          singular_vals_type = ir.F64Type.get()
          lwork = _lapack.zgesdd_work_size(
              m, n, compute_uv, full_matrices
          )
          workspace = [
              ([_lapack.gesdd_iwork_size(m, n)], i32_type),
              (
                  [_lapack.cgesdd_rwork_size(m, n, int(compute_uv))],
                  ir.F64Type.get(),
              ),
              ([lwork], a_type.element_type),
          ]
          workspace_layouts = [[0], [0], [0]]
      else:
          raise NotImplementedError(f"Unsupported dtype {dtype}")
  
      scalar_layout = []
      layout = (num_bd, num_bd + 1) + tuple(range(num_bd - 1, -1, -1))
  
      shape_type_pairs: Sequence[ShapeTypePair] = [
          (a_shape_vals, a_type.element_type),
          (batch_dims_vals + (min(m, n),), singular_vals_type),
          (
              batch_dims_vals + (m, m if full_matrices else min(m, n)),
              a_type.element_type,
          ),
          (
              batch_dims_vals + (n if full_matrices else min(m, n), n),
              a_type.element_type,
          ),
          (batch_dims_vals, i32_type),
      ] + workspace
      result_types, result_shapes = mk_result_types_and_shapes(
          shape_type_pairs
      )
      out = custom_call(
          fn,
          result_types=result_types,
          operands=[
              hlo_s32(int(full_matrices)),
              hlo_s32(int(compute_uv)),
              batch_size_val,
              hlo_s32(m),
              hlo_s32(n),
              hlo_s32(lwork),
              a,
          ],
          operand_layouts=[scalar_layout] * 6 + [layout],
          result_layouts=[
              layout,
              (num_bd,) + tuple(range(num_bd - 1, -1, -1)),
              layout,
              layout,
              tuple(range(num_bd - 1, -1, -1)),
          ]
          + workspace_layouts,
          operand_output_aliases={6: 0},
          result_shapes=result_shapes,
      ).results
      return out[1: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 6 for calling a global function with example usage: gesdd_hlo(dtype, a) and returns: out