How to use the jaxlib.cublas_kernels.build_getrf_batched_descriptor function in jaxlib

def getrf(c, a):
  """LU decomposition."""
  a_shape = c.GetShape(a)
  dtype = a_shape.element_type()
  dims = a_shape.dimensions()
  assert len(dims) >= 2
  m, n = dims[-2:]
  batch_dims = tuple(dims[:-2])
  num_bd = len(batch_dims)
  batch = _prod(batch_dims)

  if batch > 1 and m == n and m // batch <= 128:
    lwork, opaque = cublas_kernels.build_getrf_batched_descriptor(
      np.dtype(dtype), batch, m)
    workspace = _Shape.array_shape(np.dtype(np.int8), (lwork,), (0,))
    kernel = b"cublas_getrf_batched"
  else:
    lwork, opaque = cusolver_kernels.build_getrf_descriptor(
        np.dtype(dtype), batch, m, n)
    workspace = _Shape.array_shape(dtype, (lwork,), (0,))
    kernel = b"cusolver_getrf"

  out = c.CustomCall(
      kernel,
      operands=(a,),
      shape_with_layout=_Shape.tuple_shape((
          _Shape.array_shape(
              dtype, batch_dims + (m, n),
              (num_bd, num_bd + 1) + tuple(range(num_bd - 1, -1, -1))),