code: def _gesvd_hlo(
      platform,
      gpu_solver,
      have_jacobi_solver,
      dtype,
      a,
      full_matrices=True,
      compute_uv=True,
  ):
      """Singular value decomposition."""
      a_type = ir.RankedTensorType(a.type)
      dims = a_type.shape
      assert len(dims) >= 2
      m, n = dims[-2:]
      batch_dims = tuple(dims[:-2])
      num_bd = len(batch_dims)
      b = math.prod(batch_dims)
      if ir.ComplexType.isinstance(a_type.element_type):
          singular_vals_type = ir.ComplexType(
              a_type.element_type
          ).element_type
      else:
          singular_vals_type = a_type.element_type
  
      scalar_layout = tuple(range(num_bd - 1, -1, -1))
      vector_layout = (num_bd,) + tuple(range(num_bd - 1, -1, -1))
      i32_type = ir.IntegerType.get_signless(32)
  
      # NVIDIA's batched Jacobi solver supports a maximum matrix size of 32x32, but
      # the unbatched solver has no such limit. The unbatched solver appears to
      # outperform gesvd for small-moderate matrices, e.g., see:
      # https://developer.download.nvidia.com/video/gputechconf/gtc/2019/presentation/s9226-fast-singular-value-decomposition-on-gpus-v2.pdf
      # slide 5.
      if have_jacobi_solver and m <= 1024 and n <= 1024:
          # The gesvdjbatched kernel doesn't support "econ" mode. We will use that
          # kernel only if b > 1 and m <= 32 and n <= 32.
          econ = not full_matrices and (b <= 1 or m > 32 or n > 32)
          lwork, opaque = gpu_solver.build_gesvdj_descriptor(
              np.dtype(dtype), b, m, n, compute_uv, 1 if econ else 0
          )
          k = min(m, n)
          matrix_layout = (num_bd, num_bd + 1) + tuple(
              range(num_bd - 1, -1, -1)
          )
          _, s, u, v, info, _ = custom_call(
              f"{platform}solver_gesvdj",
              result_types=[
                  a.type,
                  ir.RankedTensorType.get(
                      batch_dims + (min(m, n),), singular_vals_type
                  ),
                  ir.RankedTensorType.get(
                      batch_dims + (m, k if econ else m),
                      a_type.element_type,
                  ),
                  ir.RankedTensorType.get(
                      batch_dims + (n, k if econ else n),
                      a_type.element_type,
                  ),
                  ir.RankedTensorType.get(batch_dims, i32_type),
                  ir.RankedTensorType.get([lwork], a_type.element_type),
              ],
              operands=[a],
              backend_config=opaque,
              operand_layouts=[matrix_layout],
              result_layouts=[
                  matrix_layout,
                  vector_layout,
                  matrix_layout,
                  matrix_layout,
                  scalar_layout,
                  [0],
              ],
              operand_output_aliases={0: 0},
          ).results
          vt = hlo.transpose(
              v,
              dense_int_array(
                  np.array(tuple(range(num_bd)) + (num_bd + 1, num_bd))
              ),
          )
          if np.issubdtype(dtype, np.complexfloating):
              vt = hlo.complex(hlo.real(vt), hlo.negate(hlo.imag(vt)))
          if not full_matrices and not econ:
              u = hlo.slice(
                  u,
                  dense_int_array(np.zeros([len(dims)], np.int64)),
                  dense_int_array(
                      np.array(batch_dims + (m, min(m, n)))
                  ),
                  dense_int_array(np.ones([len(dims)], np.int64)),
              )
              vt = hlo.slice(
                  vt,
                  dense_int_array(np.zeros([len(dims)], np.int64)),
                  dense_int_array(
                      np.array(batch_dims + (min(m, n), n))
                  ),
                  dense_int_array(np.ones([len(dims)], np.int64)),
              )
      elif m < n:
          lwork, opaque = gpu_solver.build_gesvd_descriptor(
              np.dtype(dtype), b, n, m, compute_uv, full_matrices
          )
          k = n if full_matrices else m
          matrix_layout = (num_bd + 1, num_bd) + tuple(
              range(num_bd - 1, -1, -1)
          )
          _, s, vt, u, info, _ = custom_call(
              f"{platform}solver_gesvd",
              result_types=[
                  a.type,
                  ir.RankedTensorType.get(
                      batch_dims + (min(m, n),), singular_vals_type
                  ),
                  ir.RankedTensorType.get(
                      batch_dims + (k, n), a_type.element_type
                  ),
                  ir.RankedTensorType.get(
                      batch_dims + (m, m), a_type.element_type
                  ),
                  ir.RankedTensorType.get(batch_dims, i32_type),
                  ir.RankedTensorType.get([lwork], a_type.element_type),
              ],
              operands=[a],
              backend_config=opaque,
              operand_layouts=[matrix_layout],
              result_layouts=[
                  matrix_layout,
                  vector_layout,
                  matrix_layout,
                  matrix_layout,
                  scalar_layout,
                  [0],
              ],
              operand_output_aliases={0: 0},
          ).results
      else:
          lwork, opaque = gpu_solver.build_gesvd_descriptor(
              np.dtype(dtype), b, m, n, compute_uv, full_matrices
          )
          k = m if full_matrices else n
          matrix_layout = (num_bd, num_bd + 1) + tuple(
              range(num_bd - 1, -1, -1)
          )
          _, s, u, vt, info, _ = custom_call(
              f"{platform}solver_gesvd",
              result_types=[
                  a.type,
                  ir.RankedTensorType.get(
                      batch_dims + (min(m, n),), singular_vals_type
                  ),
                  ir.RankedTensorType.get(
                      batch_dims + (m, k), a_type.element_type
                  ),
                  ir.RankedTensorType.get(
                      batch_dims + (n, n), a_type.element_type
                  ),
                  ir.RankedTensorType.get(batch_dims, i32_type),
                  ir.RankedTensorType.get([lwork], a_type.element_type),
              ],
              operands=[a],
              backend_config=opaque,
              operand_layouts=[matrix_layout],
              result_layouts=[
                  matrix_layout,
                  vector_layout,
                  matrix_layout,
                  matrix_layout,
                  scalar_layout,
                  [0],
              ],
              operand_output_aliases={0: 0},
          ).results
      return s, u, vt, info
  
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 8 for calling a global function with example usage: _gesvd_hlo(platform, gpu_solver, have_jacobi_solver, dtype, a, full_matrices, compute_uv) and returns: s,  u,  vt,  info