code: def _lu_pivots_to_permutation_hlo(
      platform, gpu_linalg, pivots, *, permutation_size
  ):
      """Kernel for the transformation of pivots to permutations on GPU."""
      typ = ir.RankedTensorType(pivots.type)
      dims = typ.shape
      i32_type = ir.IntegerType.get_signless(32)
  
      assert typ.element_type == i32_type, typ
  
      batch_size = _prod(dims[:-1])
      pivot_size = dims[-1]
  
      if not gpu_linalg:
          raise GpuLibNotLinkedError()
  
      opaque = gpu_linalg.lu_pivots_to_permutation_descriptor(
          batch_size, pivot_size, permutation_size
      )
      pivots_layout = tuple(range(len(dims) - 1, -1, -1))
      permutations_layout = pivots_layout
      permutations_dims = list(dims)
      permutations_dims[-1] = permutation_size
      permutations_type = ir.RankedTensorType.get(
          permutations_dims, i32_type
      )
      return custom_call(
          f"{platform}_lu_pivots_to_permutation",
          result_types=[permutations_type],
          operands=[pivots],
          backend_config=opaque,
          operand_layouts=[pivots_layout],
          result_layouts=[permutations_layout],
      ).results
  
file_path: jax/jaxlib/gpu_linalg.py
imports: import functools
  from functools import partial
  import importlib
  import operator
  import jaxlib.mlir.ir
  from .hlo_helpers import custom_call
  from .gpu_common_utils import GpuLibNotLinkedError
  from jaxlib import xla_client
label_desc: Use image node_id 1 for calling a global function with example usage: _lu_pivots_to_permutation_hlo(platform, gpu_linalg, pivots) and returns: custom_call