code: def trsm_hlo(
      dtype,
      alpha,
      a,
      b,
      left_side=False,
      lower=False,
      trans_a=False,
      conj_a=False,
      diag=False,
      *,
      b_shape_vals: tuple[DimensionSize, ...],
  ):
      _lapack.initialize()
      b_type = ir.RankedTensorType(b.type)
  
      m, n = b_shape_vals[-2:]
      batch_dims_vals = b_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)
          )
  
      if dtype == np.float32:
          fn = "blas_strsm"
      elif dtype == np.float64:
          fn = "blas_dtrsm"
      elif dtype == np.complex64:
          fn = "blas_ctrsm"
      elif dtype == np.complex128:
          fn = "blas_ztrsm"
      else:
          raise NotImplementedError(f"Unsupported dtype {dtype}")
  
      if conj_a and not trans_a:
          raise NotImplementedError(
              "Conjugation without transposition not supported"
          )
      scalar_layout = []
      layout = (num_bd, num_bd + 1) + tuple(range(num_bd - 1, -1, -1))
      result_types, result_shapes = mk_result_types_and_shapes(
          [(b_shape_vals, b_type.element_type)]
      )
      return custom_call(
          fn,
          result_types=result_types,
          operands=[
              hlo_s32(int(left_side)),
              hlo_s32(int(lower)),
              hlo_s32((2 if conj_a else 1) if trans_a else 0),
              hlo_s32(int(diag)),
              ensure_hlo_s32(m),
              ensure_hlo_s32(n),
              batch_size_val,
              alpha,
              a,
              b,
          ],
          operand_layouts=[scalar_layout] * 8 + [layout] * 2,
          result_layouts=[layout],
          operand_output_aliases={9: 0},
          result_shapes=result_shapes,
      ).results
  
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 1 for calling a global function with example usage: trsm_hlo(dtype, alpha, a, b, left_side, lower, trans_a, conj_a, diag) and returns: custom_call