code: def logit_segment(self, logits, padding_mask):
      preds = logits.argmax(dim=-1)
  
      if padding_mask.any():
          preds[padding_mask] = -1  # mark pad
      uniques = []
  
      bsz, tsz, csz = logits.shape
  
      for p in preds:
          uniques.append(
              p.cpu().unique_consecutive(
                  return_inverse=True, return_counts=True
              )
          )
  
      new_tsz = max(u[0].numel() for u in uniques)
      new_logits = logits.new_zeros(bsz, new_tsz, csz)
      new_pad = padding_mask.new_zeros(bsz, new_tsz)
  
      for b in range(bsz):
          u, idx, c = uniques[b]
          keep = u != -1
  
          if self.cfg.remove_zeros:
              keep.logical_and_(u != 0)
  
          if self.training and not self.cfg.mean_pool_join:
              u[0] = 0
              u[1:] = c.cumsum(0)[:-1]
              m = c > 1
              r = torch.rand(m.sum())
              o = (c[m] * r).long()
              u[m] += o
              new_logits[b, : u.numel()] = logits[b, u]
          else:
              new_logits[b].index_add_(
                  dim=0,
                  index=idx.to(new_logits.device),
                  source=logits[b],
              )
              new_logits[b, : c.numel()] /= c.unsqueeze(-1).to(
                  new_logits.device
              )
  
          new_sz = keep.sum()
          if not keep.all():
              kept_logits = new_logits[b, : c.numel()][keep]
              new_logits[b, :new_sz] = kept_logits
  
          if new_sz < new_tsz:
              pad = new_tsz - new_sz
              new_logits[b, -pad:] = 0
              new_pad[b, -pad:] = True
  
      return new_logits, new_pad
  
imports: from dataclasses import dataclass
  from enum import Enum, auto
  import math
  import numpy
  from typing import Tuple, List, Optional, Dict
  import torch
  import torch.nn
  import torch.nn.functional
  from torch import autograd
  from fairseq import checkpoint_utils, utils
  from fairseq.dataclass import FairseqDataclass
  from fairseq.models import BaseFairseqModel, register_model
  from fairseq.modules import SamePad, TransposeLast
label_desc: Use image node_id 1 for calling the JoinSegmenter obj's underlying member method code with example usage: obj.logit_segment(logits, padding_mask) and returns: new_logits,  new_pad