code: def make_dxx_lut(
      layout, block, step, trans, device, transform=lambda idx: idx
  ):
      # load-balancing
      _empty = torch.tensor([], dtype=torch.int64, device=layout.device)
      segments = _empty.clone()
      column = _empty.clone()
      depth = _empty.clone()
      lockid = _empty.clone()
      maxid = _empty.clone()
      offsets = _empty.clone()
      current_offset = 0
      current_maxid = 0
      for z in range(layout.size(0)):
          if trans:
              sizes = torch.sum(layout[z, :, :], 1)
          else:
              sizes = torch.sum(layout[z, :, :], 0)
          (
              z_segments,
              z_column,
              z_lockid,
              z_maxid,
              z_offsets,
          ) = _sparse_matmul.load_balance(sizes, block)
          z_depth = z * torch.ones_like(z_segments)
          z_lockid[z_lockid > 0] += current_maxid
          current_maxid = z_lockid.max()
          # concatenate depth
          segments = torch.cat((segments, z_segments))
          column = torch.cat((column, z_column))
          depth = torch.cat((depth, z_depth))
          maxid = torch.cat((maxid, z_maxid))
          offsets = torch.cat((offsets, current_offset + z_offsets))
          lockid = torch.cat((lockid, z_lockid))
          current_offset += layout[z, :, :].sum()
      segments *= step
      # pointer increments
      if trans:
          nnz = layout.nonzero()
      else:
          nnz = layout.transpose(1, 2).nonzero()
      num_blocks = nnz.size(0)
      offsets = torch.min(
          offsets, (num_blocks - 1) * torch.ones_like(offsets)
      )
      idx = transform(nnz[:, 2] * block)
      xincs = idx.clone()
      xincs[1:] -= idx[:-1]
      # divide block into multiple steps
      div = block // step
      xincs = xincs.view(-1, 1).repeat(1, div)
      xincs[:, 1:] = step
      xincs[:, 0] -= (div - 1) * step
      # first increment for each reduction is actually the offset
      xincs[offsets[segments > 0], 0] = idx[offsets[segments > 0]]
      xincs = xincs.view(-1)
      # block-mode input increments
      if trans:
          widx = torch.arange(num_blocks)
      else:
          widx = _empty.clone()
          current_offset = 0
          for z in range(layout.size(0)):
              layoutw = layout[z, :, :].clone()
              msum = layoutw.sum()
              layoutw[layoutw > 0] = 1 + torch.arange(msum)
              widx = torch.cat(
                  (widx, current_offset + layoutw.T[layoutw.T > 0] - 1)
              )
              current_offset += msum
      widx = widx
      wincs = widx * block * block
      wincs[1:] -= widx[:-1] * block * block
      wincs = wincs.view(-1, 1).repeat(1, div)
      if trans:
          wincs[:, 1:] = step
          wincs[:, 0] -= (div - 1) * step
      else:
          wincs[:, 1:] = step * block
          wincs[:, 0] -= (div - 1) * step * block
      wincs[offsets[segments > 0], 0] = widx[offsets[segments > 0]]
      wincs = wincs.view(-1)
      # adjust offset and segment size
      offsets *= 2 * div
      segments *= div
      # create header
      width = column.size(0)
      offsets += 6 * width
      header = (
          torch.stack(
              (offsets, segments, column, depth, lockid, maxid), dim=1
          )
          .view(-1)
          .contiguous()
      )
      incs = torch.stack((xincs, wincs), dim=1).view(-1).contiguous()
      incs = torch.cat(
          (incs, torch.zeros(2, device=incs.device, dtype=incs.dtype))
      )
      # create lut
      lut = torch.cat((header, incs))
      lut = lut.type(torch.int32).to(device)
      # create locks
      num_locks = max(1, lockid.max())
      return lut, num_locks, width, None
  
imports: import importlib
  import torch
  import triton
  import triton.language
  import triton._C.libtriton
  from deepspeed.accelerator import get_accelerator
label_desc: Use image node_id 5 for calling the _sparse_matmul obj's underlying member method code with example usage: obj.make_dxx_lut(layout, block, step, trans, device, transform) and returns: lut,  num_locks,  width,  None