class_docstr: Ragged unembedding implementation. This implementation will gather only the last token
  of each sequence in the ragged inflight batch and calculate the logits only for those rows.
class_docstr_tok: ["Ragged","unembedding","implementation",".","This","implementation","will","gather","only","the","last","token","of","each","sequence","in","the","ragged","inflight","batch","and","calculate","the","logits","only","for","those","rows","."]
code: def output(self) -> torch.Tensor:
      return self._output
  
file_path: DeepSpeed/deepspeed/inference/v2/modules/implementations/unembed/ragged_unembed.py
imports: from typing import Any, Dict, Optional
  import torch
  from deepspeed.accelerator import get_accelerator
  from ....allocator import empty_from
  from ....inference_utils import DtypeEnum, ActivationType
  from ....kernels.core_ops import CUDAFPLN, BlasLibLinear, CUDARMSNorm, CUDABiasActivation
  from ....kernels.ragged_ops import RaggedLogitsGather
  from ....ragged import RaggedBatchWrapper
  from ...interfaces import DSUnembedBase, DSUnembedRegistry
  from ...configs import DSUnembedConfig
label_desc: Use image node_id 4 for calling the DSRaggedUnembed obj's underlying member method code with example usage: obj.output() and returns: self