code: def logit_segment(
      self,
      logits: torch.Tensor,
      padding_mask: torch.Tensor,
  ) -> torch.Tensor:
      assert check_argument_types()
      preds = logits.argmax(dim=-1)
  
      if padding_mask.any():
          preds[padding_mask] = -1  # mark pad
      uniques = []
  
      batch_size, time_length, channel_size = logits.shape
  
      for p in preds:
          uniques.append(
              p.cpu().unique_consecutive(
                  return_inverse=True, return_counts=True
              )
          )
  
      new_time_length = max(u[0].numel() for u in uniques)
      new_logits = logits.new_zeros(
          batch_size, new_time_length, channel_size
      )
      new_pad = padding_mask.new_zeros(batch_size, new_time_length)
  
      for b in range(batch_size):
          value, index, count = uniques[b]
          keep = value != -1
  
          if self.remove_zeros:
              keep.logical_and_(value != 0)
  
          if self.training and not self.mean_pool_join:
              value[0] = 0
              value[1:] = count.cumsum(0)[:-1]
              part = count > 1
              random = torch.rand(part.sum())
              value[part] += (count[part] * random).long()
              new_logits[b, : value.numel()] = logits[b, value]
          else:
              new_logits[b].index_add_(
                  dim=0,
                  index=index.to(new_logits.device),
                  source=logits[b],
              )
              new_logits[b, : count.numel()] = new_logits[
                  b, : count.numel()
              ] / count.unsqueeze(-1).to(new_logits.device)
  
          new_size = keep.sum()
          if not keep.all():
              kept_logits = new_logits[b, : count.numel()][keep]
              new_logits[b, :new_size] = kept_logits
  
          if new_size < new_time_length:
              pad = new_time_length - new_size
              new_logits[b, -pad:] = 0
              new_pad[b, -pad:] = True
      return new_logits, new_pad
  
file_path: espnet/espnet2/uasr/segmenter/join_segmenter.py
imports: import argparse
  from typing import Dict, Optional
  import torch
  from typeguard import check_argument_types
  from espnet2.uasr.segmenter.abs_segmenter import AbsSegmenter
  from espnet2.utils.types import str2bool
label_desc: Use image node_id 3 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